aimbat.core

def add_data_to_project(
    session: Session,
    data_sources: Sequence[str | os.PathLike],
    data_type: DataType,
    station_id: uuid.UUID | None = None,
    event_id: uuid.UUID | None = None,
    dry_run: bool = False,
    disable_progress_bar: bool = True,
) -> None:
    """Add data sources to the AIMBAT database.

    What gets created depends on which capabilities `data_type` supports:

    - Station + event + seismogram: all three records are created and linked,
      and an `AimbatDataSource` entry is stored.
    - Station or event only (e.g. `JSON_STATION`, `JSON_EVENT`): only the
      relevant metadata records are created; no seismogram or data source entry
      is stored.

    Use `station_id` or `event_id` to skip extracting station or event metadata
    from the data source and link to a pre-existing record instead.

    Args:
        session: The SQLModel database session.
        data_sources: List of data sources to add.
        data_type: Type of data.
        station_id: UUID of an existing station to use instead of extracting
            one from each data source.
        event_id: UUID of an existing event to use instead of extracting one
            from each data source.
        dry_run: If True, do not commit changes to the database.
        disable_progress_bar: Do not display progress bar.
    """

    logger.info(f"Adding {len(data_sources)} {data_type} data sources to project.")

    if station_id is not None and session.get(AimbatStation, station_id) is None:
        raise ValueError(f"No station found with ID {station_id}.")
    if event_id is not None and session.get(AimbatEvent, event_id) is None:
        raise ValueError(f"No event found with ID {event_id}.")

    # Snapshot existing IDs before entering the savepoint so we can identify
    # what would be new vs reused when running a dry run.
    if dry_run:
        existing_station_ids = set(session.exec(select(AimbatStation.id)).all())
        existing_event_ids = set(session.exec(select(AimbatEvent.id)).all())
        existing_seismogram_ids = set(session.exec(select(AimbatSeismogram.id)).all())

    try:
        added_datasources: list[AimbatDataSource] = []
        with session.begin_nested() as nested:
            for datasource in track(
                sequence=data_sources,
                description="Adding data ...",
                disable=disable_progress_bar,
            ):
                result = _process_datasource(
                    session, datasource, data_type, station_id, event_id
                )
                if result is not None:
                    added_datasources.append(result)

            if dry_run:
                logger.info("Dry run: displaying data that would be added.")
                if added_datasources:
                    session.flush()
                    _print_dry_run_results(
                        added_datasources,
                        existing_station_ids,
                        existing_event_ids,
                        existing_seismogram_ids,
                    )
                nested.rollback()
                logger.info("Dry run complete. Rolling back changes.")
                return

        session.commit()
        logger.info("Data added successfully.")

    except Exception as e:
        logger.error(f"Failed to add data. Rolling back changes. Error: {e}")
        raise

def build_iccs_from_snapshot(session: Session, snapshot_id: UUID) -> BoundICCS:
    """Build a read-only BoundICCS from a snapshot's parameters and live waveform data.

    Uses the snapshot's event and seismogram parameters (window, t1, flip, select,
    bandpass, etc.) but reads waveform data from the live datasources. Seismograms
    added after the snapshot was taken are not included in the snapshot — their live
    parameters are used instead. No DB writes occur at any point.

    Args:
        session: Database session.
        snapshot_id: ID of the AimbatSnapshot to load.

    Returns:
        BoundICCS instance built from the snapshot parameters.

    Raises:
        ValueError: If no snapshot with the given ID is found.
    """
    snapshot = session.get(AimbatSnapshot, snapshot_id)
    if snapshot is None:
        raise ValueError(f"Snapshot {snapshot_id} not found.")

    ep = snapshot.event_parameters_snapshot
    snap_params = AimbatEventParametersBase.model_validate(ep)

    # Build a map from seismogram_parameters_id → snapshot parameters
    snap_seis_map = {
        sp.seismogram_parameters_id: sp
        for sp in snapshot.seismogram_parameters_snapshots
    }

    seismograms = []
    for seis in snapshot.event.seismograms:
        snap_sp = snap_seis_map.get(seis.parameters.id)
        if snap_sp is None:
            # Seismogram was added after the snapshot — use live parameters
            seis_params = AimbatSeismogramParametersBase.model_validate(seis.parameters)
        else:
            seis_params = AimbatSeismogramParametersBase.model_validate(snap_sp)
        seismograms.append(
            MiniICCSSeismogram(
                begin_time=seis.begin_time,
                delta=seis.delta,
                data=seis.data,
                t0=seis.t0,
                t1=seis_params.t1,
                flip=seis_params.flip,
                select=seis_params.select,
                extra={"id": seis.id},
            )
        )

    iccs = ICCS(
        seismograms=seismograms,
        window_pre=snap_params.window_pre,
        window_post=snap_params.window_post,
        bandpass_apply=snap_params.bandpass_apply,
        bandpass_fmin=snap_params.bandpass_fmin,
        bandpass_fmax=snap_params.bandpass_fmax,
        min_ccnorm=snap_params.min_ccnorm,
        context_width=settings.context_width,
    )
    return BoundICCS(
        iccs=iccs,
        event_id=snapshot.event_id,
        created_at=Timestamp.now("UTC"),
    )

def clear_iccs_cache() -> None:
    """Clear the process-level ICCS cache."""
    _iccs_cache.clear()

def create_iccs_instance(session: Session, event: AimbatEvent) -> BoundICCS:
    """Return a BoundICCS instance for the given event.

    Returns the cached instance when it is still fresh (i.e. `event.last_modified`
    has not advanced since the instance was created). Otherwise builds a new one
    and updates the cache.

    `MiniICCSSeismogram` instances are constructed directly from each
    `AimbatSeismogram`, passing `data` by reference to the read-only io cache.
    No waveform data is copied. The session does not need to remain open after
    this call.

    Args:
        session: Database session.
        event: AimbatEvent.

    Returns:
        BoundICCS instance tied to the given event.
    """
    cached = _iccs_cache.get(event.id)
    if cached is not None and not cached.is_stale(event):
        logger.debug(f"Returning cached BoundICCS for event {event.id}.")
        return cached

    logger.info(f"Creating ICCS instance for event {event.id}.")
    bound = BoundICCS(
        iccs=_build_iccs(event),
        event_id=event.id,
        created_at=Timestamp.now("UTC"),
    )
    _iccs_cache[event.id] = bound
    return bound

def create_project(engine: Engine) -> None:
    """Initializes a new AIMBAT project database schema and triggers.

    Args:
        engine: The SQLAlchemy/SQLModel Engine instance connected to the target database.

    Raises:
        RuntimeError: If a project schema already exists in the target database.
    """

    # Import locally to ensure SQLModel registers all table metadata before create_all()
    import aimbat.models  # noqa: F401

    logger.info(f"Creating new project in {engine.url}")

    if _project_exists(engine):
        raise RuntimeError(
            f"Unable to create a new project: project already exists at {engine.url}!"
        )

    logger.debug("Creating database tables and loading defaults.")

    SQLModel.metadata.create_all(engine)

    if engine.name == "sqlite":
        with engine.begin() as connection:
            # Trigger 1: Handle updates to existing rows
            connection.execute(
                text("""
                CREATE TRIGGER IF NOT EXISTS single_default_event_update
                BEFORE UPDATE ON aimbatevent
                FOR EACH ROW WHEN NEW.is_default = TRUE
                BEGIN
                    UPDATE aimbatevent SET is_default = NULL 
                    WHERE is_default = TRUE AND id != NEW.id;
                END;
            """)
            )

            # Trigger 2: Handle brand new default events being inserted
            connection.execute(
                text("""
                CREATE TRIGGER IF NOT EXISTS single_default_event_insert
                BEFORE INSERT ON aimbatevent
                FOR EACH ROW WHEN NEW.is_default = TRUE
                BEGIN
                    UPDATE aimbatevent SET is_default = NULL
                    WHERE is_default = TRUE;
                END;
            """)
            )

            # Trigger 3: Track last modification time when event parameters change
            connection.execute(
                text("""
                CREATE TRIGGER IF NOT EXISTS event_modified_on_params_update
                AFTER UPDATE ON aimbateventparameters
                BEGIN
                    UPDATE aimbatevent SET last_modified = datetime('now')
                    WHERE id = NEW.event_id;
                END;
            """)
            )

            # Trigger 4: Track last modification time when seismogram parameters change
            connection.execute(
                text("""
                CREATE TRIGGER IF NOT EXISTS event_modified_on_seis_params_update
                AFTER UPDATE ON aimbatseismogramparameters
                BEGIN
                    UPDATE aimbatevent 
                    SET last_modified = strftime('%Y-%m-%d %H:%M:%f', 'now')
                    WHERE id = (
                        SELECT event_id FROM aimbatseismogram
                        WHERE id = NEW.seismogram_id
                    );
                END;
            """)
            )

def create_snapshot(
    session: Session, event: AimbatEvent, comment: str | None = None
) -> None:
    """Create a snapshot of the AIMBAT processing parameters.

    Args:
        session: Database session.
        event: AimbatEvent.
        comment: Optional comment.
    """

    logger.info(f"Creating snapshot for event with id={event.id} with {comment=}.")

    event_parameters_snapshot = AimbatEventParametersSnapshot.model_validate(
        event.parameters,
        update={
            "id": uuid.uuid4(),  # we don't want to carry over the id from the input event parameters
            "parameters_id": event.parameters.id,
        },
    )
    logger.debug(
        f"Adding event parameters snapshot with id={event_parameters_snapshot.id} to snapshot."
    )

    seismogram_parameter_snapshots = []
    for aimbat_seismogram in event.seismograms:
        seismogram_parameter_snapshot = AimbatSeismogramParametersSnapshot.model_validate(
            aimbat_seismogram.parameters,
            update={
                "id": uuid.uuid4(),  # we don't want to carry over the id from the input seismogram parameters
                "seismogram_parameters_id": aimbat_seismogram.parameters.id,
            },
        )
        logger.debug(
            f"Adding seismogram parameters snapshot with id={seismogram_parameter_snapshot.id} to snapshot."
        )
        seismogram_parameter_snapshots.append(seismogram_parameter_snapshot)

    aimbat_snapshot = AimbatSnapshot(
        event=event,
        event_parameters_snapshot=event_parameters_snapshot,
        seismogram_parameters_snapshots=seismogram_parameter_snapshots,
        comment=comment,
    )
    session.add(aimbat_snapshot)
    session.commit()

def delete_event(session: Session, event: AimbatEvent) -> None:
    """Delete an AimbatEvent from the database.

    Args:
        session: Database session.
        event: Event to delete.
    """

    logger.info(f"Deleting event {event.id}.")

    session.delete(event)
    session.commit()

def delete_event_by_id(session: Session, event_id: UUID) -> None:
    """Delete an AimbatEvent from the database by ID.

    Args:
        session: Database session.
        event_id: Event ID.

    Raises:
        NoResultFound: If no AimbatEvent is found with the given ID.
    """

    logger.debug(f"Getting event with id={event_id}.")

    event = session.get(AimbatEvent, event_id)
    if event is None:
        raise NoResultFound(f"Unable to find event using id: {event_id}.")
    delete_event(session, event)

def delete_project(engine: Engine) -> None:
    """Delete the AIMBAT project.

    Raises:
        RuntimeError: If unable to delete project.
    """

    logger.info(f"Deleting project in {engine=}.")

    if _project_exists(engine):
        if engine.driver == "pysqlite":
            database = engine.url.database
            engine.dispose()
            if database == ":memory:":
                logger.info("Running database in memory, nothing to delete.")
                return
            elif database:
                project_path = Path(database)
                logger.info(f"Deleting project file: {project_path=}")
                project_path.unlink()
                return
    raise RuntimeError("Unable to find/delete project.")

def delete_seismogram(session: Session, seismogram: AimbatSeismogram) -> None:
    """Delete an AimbatSeismogram from the database.

    Args:
        session: Database session.
        seismogram: Seismogram to delete.
    """

    logger.info(f"Deleting seismogram {seismogram.id}.")

    session.delete(seismogram)
    session.commit()

def delete_seismogram_by_id(session: Session, seismogram_id: uuid.UUID) -> None:
    """Delete an AimbatSeismogram from the database by ID.

    Args:
        session: Database session.
        seismogram_id: Seismogram ID.

    Raises:
        NoResultFound: If no AimbatSeismogram is found with the given ID.
    """

    logger.debug(f"Getting seismogram with id={seismogram_id}.")

    seismogram = session.get(AimbatSeismogram, seismogram_id)
    if seismogram is None:
        raise NoResultFound(f"No AimbatSeismogram found with {seismogram_id=}")
    delete_seismogram(session, seismogram)

def delete_snapshot(session: Session, snapshot: AimbatSnapshot) -> None:
    """Delete an AIMBAT parameter snapshot.

    Args:
        session: Database session.
        snapshot: Snapshot.
    """

    logger.info(f"Deleting snapshot {snapshot.id}.")

    session.delete(snapshot)
    session.commit()

def delete_snapshot_by_id(session: Session, snapshot_id: uuid.UUID) -> None:
    """Delete an AIMBAT parameter snapshot.

    Args:
        session: Database session.
        snapshot_id: Snapshot id.
    """

    logger.debug(f"Searching for snapshot with id {snapshot_id}.")

    snapshot = session.get(AimbatSnapshot, snapshot_id)

    if snapshot is None:
        raise ValueError(
            f"Unable to delete snapshot: snapshot with id={snapshot_id} not found."
        )

    delete_snapshot(session, snapshot)

def delete_station(session: Session, station: AimbatStation) -> None:
    """Delete an AimbatStation from the database.

    Args:
        session: Database session.
        station: Station to delete.
    """

    logger.info(f"Deleting station {station.id}.")

    session.delete(station)
    session.commit()

def delete_station_by_id(session: Session, station_id: uuid.UUID) -> None:
    """Delete an AimbatStation from the database by ID.

    Args:
        session: Database session.
        station_id: Station ID.

    Raises:
        NoResultFound: If no AimbatStation is found with the given ID.
    """

    logger.debug(f"Getting station with id={station_id}.")

    station = session.get(AimbatStation, station_id)
    if station is None:
        raise NoResultFound(f"No AimbatStation found with {station_id=}")
    delete_station(session, station)

def dump_data_table_to_json(session: Session) -> str:
    """Dump the table data to json."""

    logger.info("Dumping AIMBAT datasources table to json.")

    adapter: TypeAdapter[Sequence[AimbatDataSource]] = TypeAdapter(
        Sequence[AimbatDataSource]
    )
    aimbat_datasource = session.exec(select(AimbatDataSource)).all()
    return adapter.dump_json(aimbat_datasource).decode("utf-8")

def dump_event_parameter_table_to_json(
    session: Session,
    all_events: bool,
    as_string: bool,
    event: AimbatEvent | None = None,
) -> str | dict[str, Any] | list[dict[str, Any]]:
    """Dump the event parameter table data to json.

    Args:
        session: Database session.
        all_events: Include event parameter table data for all events.
        as_string: Whether to return the result as a string.
        event: Event to dump parameter data for (only used when all_events is False).
    """

    logger.info("Dumping AIMBAT event parameter table to json.")

    if all_events:
        adapter: TypeAdapter[Sequence[AimbatEventParameters]] = TypeAdapter(
            Sequence[AimbatEventParameters]
        )
        parameters = session.exec(select(AimbatEventParameters)).all()
        if as_string:
            return adapter.dump_json(parameters).decode("utf-8")
        else:
            return adapter.dump_python(parameters, mode="json")

    if event is None:
        raise ValueError("An event must be provided when all_events is False.")

    if as_string:
        return event.parameters.model_dump_json()
    return event.parameters.model_dump(mode="json")

def dump_event_table_to_json(
    session: Session, as_string: bool = True
) -> str | list[dict[str, Any]]:
    """Dump the table data to json."""

    logger.info("Dumping AIMBAT event table to json.")
    events = session.exec(select(AimbatEvent)).all()
    event_reads = [_AimbatEventRead.from_event(e) for e in events]
    adapter: TypeAdapter[Sequence[_AimbatEventRead]] = TypeAdapter(
        Sequence[_AimbatEventRead]
    )
    if as_string:
        return adapter.dump_json(event_reads).decode("utf-8")
    return adapter.dump_python(event_reads, mode="json")

def dump_seismogram_parameter_table_to_json(
    session: Session,
    all_events: bool,
    as_string: bool,
    event: AimbatEvent | None = None,
) -> str | list[dict[str, Any]]:
    """Dump the seismogram parameter table data to json.

    Args:
        session: Database session.
        all_events: Include parameters for all events.
        as_string: Return as JSON string.
        event: Event to dump parameters for (only used when all_events is False).
    """

    logger.info("Dumping AimbatSeismogramParameters table to json.")

    adapter: TypeAdapter[Sequence[AimbatSeismogramParameters]] = TypeAdapter(
        Sequence[AimbatSeismogramParameters]
    )

    if all_events:
        parameters = session.exec(select(AimbatSeismogramParameters)).all()
    else:
        if event is None:
            raise ValueError("An event must be provided when all_events is False.")
        parameters = session.exec(
            select(AimbatSeismogramParameters)
            .join(AimbatSeismogram)
            .where(AimbatSeismogram.event_id == event.id)
        ).all()

    if as_string:
        return adapter.dump_json(parameters).decode("utf-8")
    return adapter.dump_python(parameters, mode="json")

def dump_seismogram_table_to_json(session: Session) -> str:
    """Create a JSON string from the AimbatSeismogram table data."""

    logger.info("Dumping AIMBAT seismogram table to json.")
    adapter: TypeAdapter[Sequence[AimbatSeismogram]] = TypeAdapter(
        Sequence[AimbatSeismogram]
    )
    aimbat_seismograms = session.exec(select(AimbatSeismogram)).all()

    return adapter.dump_json(aimbat_seismograms).decode("utf-8")

def dump_snapshot_tables_to_json(
    session: Session,
    all_events: bool,
    as_string: bool,
    event: AimbatEvent | None = None,
) -> str | dict[str, list[dict[str, Any]]]:
    """Dump snapshot data as a dict of lists of dicts.

    Returns a structure with three keys:

    - ``snapshots``: flat list of snapshot metadata.
    - ``event_parameters``: flat list of event parameter snapshots.
    - ``seismogram_parameters``: flat list of seismogram parameter snapshots.

    Each entry includes a ``snapshot_id`` for cross-referencing.

    Args:
        session: Database session.
        all_events: Include snapshots for all events.
        as_string: Return a JSON string when True, otherwise a dict.
        event: Event to dump snapshots for (only used when all_events is False).
    """
    logger.info(f"Dumping AimbatSnapshot tables to json with {all_events=}.")

    snapshots = get_snapshots(session, event=event, all_events=all_events)

    snapshot_adapter: TypeAdapter[Sequence[_AimbatSnapshotRead]] = TypeAdapter(
        Sequence[_AimbatSnapshotRead]
    )
    event_params_adapter: TypeAdapter[Sequence[AimbatEventParametersSnapshot]] = (
        TypeAdapter(Sequence[AimbatEventParametersSnapshot])
    )
    seis_params_adapter: TypeAdapter[Sequence[AimbatSeismogramParametersSnapshot]] = (
        TypeAdapter(Sequence[AimbatSeismogramParametersSnapshot])
    )

    snapshot_reads = [_AimbatSnapshotRead.from_snapshot(s) for s in snapshots]
    event_params = [s.event_parameters_snapshot for s in snapshots]
    seis_params = [sp for s in snapshots for sp in s.seismogram_parameters_snapshots]

    data: dict[str, list[dict[str, Any]]] = {
        "snapshots": snapshot_adapter.dump_python(snapshot_reads, mode="json"),
        "event_parameters": event_params_adapter.dump_python(event_params, mode="json"),
        "seismogram_parameters": seis_params_adapter.dump_python(
            seis_params, mode="json"
        ),
    }

    return json.dumps(data) if as_string else data

def dump_station_table_to_json(session: Session) -> str:
    """Create a JSON string from the AimbatStation table data."""

    logger.info("Dumping AIMBAT station table to json.")

    adapter: TypeAdapter[Sequence[AimbatStation]] = TypeAdapter(Sequence[AimbatStation])
    aimbat_station = session.exec(select(AimbatStation)).all()
    return adapter.dump_json(aimbat_station).decode("utf-8")

def dump_station_table_with_counts(session: Session) -> list[dict[str, Any]]:
    """Dump station table with associated seismogram and event counts to a list of dicts.

    Each dict represents a station and includes additional keys for the
    seismogram and event counts.

    Args:
        session: Database session.

    Returns: A list of dictionaries representing the stations with counts.
    """
    results = get_stations_with_event_and_seismogram_count(session)
    formatted_results = []

    for row in results:
        station_dict = row[0].model_dump(mode="json")
        station_dict["seismogram_count"] = row[1]
        station_dict["event_count"] = row[2]
        formatted_results.append(station_dict)

    return formatted_results

def get_completed_events(session: Session) -> Sequence[AimbatEvent]:
    """Get the events marked as completed.

    Args:
        session: SQL session.
    """

    statement = (
        select(AimbatEvent)
        .join(AimbatEventParameters)
        .where(AimbatEventParameters.completed == 1)
    )

    return session.exec(statement).all()