aimbat.core

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

    Args:
        session: The SQLModel database session.
        data_sources: List of data sources to add.
        data_type: Type of data.
        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.")

    # 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,
            ):
                added_datasources.append(
                    _add_datasource(session, datasource, data_type)
                )

            if dry_run:
                logger.info("Dry run: displaying data that would be added.")
                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 create_iccs_instance(session: Session) -> ICCS:
    """Create an ICCS instance for the active event.

    Args:
        session: Database session.

    Returns:
        ICCS instance.
    """

    logger.info("Creating ICCS instance for active event.")

    active_event = get_active_event(session)

    return ICCS(
        seismograms=active_event.seismograms,
        window_pre=active_event.parameters.window_pre,
        window_post=active_event.parameters.window_post,
        bandpass_apply=active_event.parameters.bandpass_apply,
        bandpass_fmin=active_event.parameters.bandpass_fmin,
        bandpass_fmax=active_event.parameters.bandpass_fmax,
        min_ccnorm=active_event.parameters.min_ccnorm,
        context_width=settings.context_width,
    )

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_active_event_update
                BEFORE UPDATE ON aimbatevent
                FOR EACH ROW WHEN NEW.active = TRUE
                BEGIN
                    UPDATE aimbatevent SET active = NULL 
                    WHERE active = TRUE AND id != NEW.id;
                END;
            """))

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

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

    Args:
        session: Database session.
        comment: Optional comment.
    """
    active_aimbat_event = get_active_event(session)

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

    event_parameters_snapshot = AimbatEventParametersSnapshot.model_validate(
        active_aimbat_event.parameters,
        update={
            "id": uuid.uuid4(),  # we don't want to carry over the id from the active parameters
            "parameters_id": active_aimbat_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 active_aimbat_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 active 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=active_aimbat_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}. {HINTS.LIST_EVENTS}"
        )
    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
) -> str | dict[str, Any] | list[dict[str, Any]]:
    """Dump the event parameter table data to json."""

    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")

    active_event = get_active_event(session)

    if as_string:
        return active_event.parameters.model_dump_json()
    return active_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
) -> str | list[dict[str, Any]]:
    """Dump the seismogram parameter table data to json."""

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

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

    if all_events:
        parameters = session.exec(select(AimbatSeismogramParameters)).all()
    else:
        parameters = session.exec(
            select(AimbatSeismogramParameters)
            .join(AimbatSeismogram)
            .join(AimbatEvent)
            .where(AimbatEvent.active == 1)
        ).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
) -> 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.
    """
    logger.info(f"Dumping AimbatSnapshot tables to json with {all_events=}.")

    snapshots = get_snapshots(session, 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 get_active_event(session: Session) -> AimbatEvent:
    """
    Return the currently active event (i.e. the one being processed).

    Args:
        session: SQL session.

    Returns:
        Active Event

    Raises
        NoResultFound: When no event is active.
    """

    logger.debug("Attempting to determine active event.")

    select_active_event = select(AimbatEvent).where(AimbatEvent.active == 1)

    # NOTE: While there technically can be no active event in the database,
    # we typically don't really want to go beyond this point when that is the
    # case. Hence we call `one` rather than `one_or_none`.
    try:
        active_event = session.exec(select_active_event).one()
    except NoResultFound:
        raise NoResultFound(f"No active event found. {HINTS.ACTIVATE_EVENT}")

    logger.debug(f"Active event: {active_event.id}")

    return active_event

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

    Args:
        session: SQL session.
    """

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

    return session.exec(select_completed_events).all()

def get_data_for_active_event(session: Session) -> Sequence[AimbatDataSource]:
    """Returns the data sources belonging to the active event.

    Args:
        session: Database session.

    Returns:
        Sequence of AimbatDataSource objects belonging to the active event.
    """

    logger.info("Getting data sources for active event.")

    statement = (
        select(AimbatDataSource)
        .join(AimbatSeismogram)
        .join(AimbatEvent)
        .where(AimbatEvent.active == 1)
    )
    return session.exec(statement).all()

def get_event_parameter(
    session: Session, name: EventParameter
) -> Timedelta | bool | float:
    """Get event parameter value for the active event.

    Args:
        session: Database session.
        name: Name of the parameter.
    """

    active_event = get_active_event(session)

    logger.info(f"Getting {name=} value for {active_event=}.")

    return getattr(active_event.parameters, name)

def get_events_using_station(
    session: Session, station: AimbatStation
) -> Sequence[AimbatEvent]:
    """Get all events that use a particular station.

    Args:
        session: Database session.
        station: Station to return events for.

    Returns: Events that use the station.
    """

    logger.info(f"Getting events for station: {station.id}.")

    select_events = (
        select(AimbatEvent)
        .join(AimbatSeismogram)
        .join(AimbatStation)
        .where(AimbatStation.id == station.id)
    )

    events = session.exec(select_events).all()

    logger.debug(f"Found {len(events)}.")

    return events

def get_seismogram_parameter(
    seismogram: AimbatSeismogram, name: SeismogramParameter
) -> bool | Timestamp:
    """Get parameter value from an AimbatSeismogram instance.

    Args:
        seismogram: Seismogram.
        name: Name of the parameter value to return.

    Returns:
        Seismogram parameter value.
    """

    logger.info(f"Getting seismogram parameter {name=} value for {seismogram=}.")

    return getattr(seismogram.parameters, name)

def get_seismogram_parameter_by_id(
    session: Session, seismogram_id: uuid.UUID, name: SeismogramParameter
) -> bool | Timestamp:
    """Get parameter value from an AimbatSeismogram by ID.

    Args:
        session: Database session.
        seismogram_id: Seismogram ID.
        name: Name of the parameter value to return.

    Returns:
        Seismogram parameter value.

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

    logger.info(f"Getting seismogram {name=} for seismogram with id={seismogram_id}.")

    aimbat_seismogram = session.get(AimbatSeismogram, seismogram_id)

    if aimbat_seismogram is None:
        raise ValueError(f"No AimbatSeismogram found with {seismogram_id=}")

    return get_seismogram_parameter(aimbat_seismogram, name)

def get_selected_seismograms(
    session: Session, all_events: bool = False
) -> Sequence[AimbatSeismogram]:
    """Get the selected seismograms for the active avent.

    Args:
        session: Database session.
        all_events: Get the selected seismograms for all events.

    Returns: Selected seismograms.
    """

    logger.info("Getting selected AIMBAT seismograms.")

    if all_events is True:
        logger.debug("Selecting seismograms for all events.")
        select_seismograms = (
            select(AimbatSeismogram)
            .join(AimbatSeismogramParameters)
            .where(AimbatSeismogramParameters.select == 1)
        )
    else:
        logger.debug("Selecting seismograms for active event only.")
        select_seismograms = (
            select(AimbatSeismogram)
            .join(AimbatSeismogramParameters)
            .join(AimbatEvent)
            .where(AimbatSeismogramParameters.select == 1)
            .where(AimbatEvent.active == 1)
        )

    seismograms = session.exec(select_seismograms).all()

    logger.debug(f"Found {len(seismograms)} selected seismograms.")

    return seismograms

def get_snapshots(
    session: Session, all_events: bool = False
) -> Sequence[AimbatSnapshot]:
    """Get the snapshots for the active avent.

    Args:
        session: Database session.
        all_events: Get the selected snapshots for all events.

    Returns: Snapshots.
    """

    logger.info("Getting AIMBAT snapshots.")

    statement = (
        select(AimbatSnapshot)
        .join(AimbatEvent)
        .where(AimbatEvent.active == True if not all_events else true())  # noqa: E712
    )

    logger.debug(f"Executing statement to get snapshots: {statement}")
    return session.exec(statement).all()