aimbat._cli

@app.command(name="iccs")
@simple_exception
def cli_iccs_run(
    *,
    autoflip: bool = False,
    autoselect: bool = False,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Run the ICCS algorithm to align seismograms for the default event.

    Iteratively cross-correlates seismograms against a running stack to refine
    arrival time picks (`t1`). If `t1` is not yet set, `t0` is used as the
    starting point.

    Args:
        autoflip: Whether to automatically flip seismograms (multiply data by -1)
            when the cross-correlation is negative.
        autoselect: Whether to automatically de-select seismograms whose
            cross-correlation with the stack falls below `min_ccnorm`.
    """
    from aimbat.db import engine
    from aimbat.core import create_iccs_instance, run_iccs, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        iccs = create_iccs_instance(session, event).iccs
        run_iccs(session, iccs, autoflip, autoselect)

@app.command(name="mccc")
@simple_exception
def cli_mccc_run(
    *,
    all_seismograms: Annotated[bool, Parameter(name="all")] = False,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Run the MCCC algorithm to refine arrival time picks for the default event.

    Multi-channel cross-correlation simultaneously determines the optimal time
    shifts for all seismograms. Results are stored in `t1`.

    Args:
        all_seismograms: Include all seismograms in MCCC processing, not just
            the currently selected ones.
    """
    from aimbat.db import engine
    from aimbat.core import create_iccs_instance, run_mccc, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        iccs = create_iccs_instance(session, event).iccs
        run_mccc(session, event, iccs, all_seismograms)

def id_parameter(model_class: type) -> Parameter:
    """Create a Parameter for a record ID with automatic UUID prefix resolution."""
    return Parameter(
        name="id",
        help="Full UUID or any unique prefix as shown in the table.",
        converter=_make_uuid_converter(model_class),
    )

def simple_exception[F: Callable[..., Any]](func: F) -> F:
    """Decorator to handle exceptions and print them to the console.

    Using this decorator prints only the exception to the console without
    traceback, and then exits. In debugging mode this decorator returns the
    callable unchanged.
    """
    from functools import wraps
    from rich.console import Console
    from rich.panel import Panel
    import sys

    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        if settings.log_level in ("TRACE", "DEBUG"):
            return func(*args, **kwargs)
        try:
            return func(*args, **kwargs)
        except Exception as e:
            console = Console()
            panel = Panel(
                f"{e}",
                title="Error",
                title_align="left",
                border_style="red",
                expand=True,
            )
            console.print(panel)
            sys.exit(1)

    return wrapper  # type: ignore

def use_event_parameter(model_class: type) -> Parameter:
    """Create a Parameter for --use-event with automatic UUID prefix resolution."""
    return Parameter(
        name="use-event",
        help="UUID (or unique prefix) of an existing event to link to instead of"
        " extracting one from each data source.",
        converter=_make_uuid_converter(model_class),
    )

def use_station_parameter(model_class: type) -> Parameter:
    """Create a Parameter for --use-station with automatic UUID prefix resolution."""
    return Parameter(
        name="use-station",
        help="UUID (or unique prefix) of an existing station to link to instead of"
        " extracting one from each data source.",
        converter=_make_uuid_converter(model_class),
    )

@app.command(name="add")
@simple_exception
def cli_data_add(
    data_sources: Annotated[
        list[Path],
        Parameter(
            name="sources",
            consume_multiple=True,
            validator=validators.Path(exists=True),
        ),
    ],
    *,
    data_type: Annotated[DataType, Parameter(name="type")] = DataType.SAC,
    station_id: Annotated[
        uuid.UUID | None, use_station_parameter(AimbatStation)
    ] = None,
    event_id: Annotated[uuid.UUID | None, use_event_parameter(AimbatEvent)] = None,
    dry_run: Annotated[bool, Parameter(name="dry-run")] = False,
    show_progress_bar: Annotated[bool, Parameter(name="progress")] = True,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Add or update data sources in the AIMBAT project.

    Each data source is processed according to `--type`. For `sac` (the
    default), AIMBAT extracts station, event, and seismogram metadata directly
    from the file. For types that cannot extract a station or event (e.g. a
    format that only carries waveform data), supply `--use-station` and/or
    `--use-event` to link to records that already exist in the project.

    Station and event deduplication is automatic: if a matching record already
    exists it is reused. Re-running `data add` on the same files is safe.

    Use `--dry-run` to preview what would be added without touching the
    database.

    Args:
        data_sources: One or more data source paths to add.
        data_type: Format of the data sources. Determines which metadata
            (station, event, seismogram) can be extracted automatically.
        dry_run: Preview which records would be added without modifying the
            database.
        show_progress_bar: Display a progress bar while ingesting sources.
    """
    from aimbat.db import engine
    from aimbat.core import add_data_to_project

    disable_progress_bar = not show_progress_bar

    with Session(engine) as session:
        add_data_to_project(
            session,
            data_sources,
            data_type,
            station_id=station_id,
            event_id=event_id,
            dry_run=dry_run,
            disable_progress_bar=disable_progress_bar,
        )

@app.command(name="dump")
@simple_exception
def cli_data_dump(
    *,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Dump the contents of the AIMBAT data source table to JSON.

    Output can be piped or redirected for use in external tools or scripts.
    """
    from aimbat.db import engine
    from aimbat.core import dump_data_table_to_json
    from rich import print_json

    with Session(engine) as session:
        print_json(dump_data_table_to_json(session))

@app.command(name="list")
@simple_exception
def cli_data_list(
    *,
    all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False,
    table_parameters: TableParameters = TableParameters(),
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Print a table of data sources registered in the AIMBAT project."""
    from aimbat.db import engine
    from aimbat.core import resolve_event, get_data_for_event
    from aimbat.utils import uuid_shortener, make_table, TABLE_STYLING
    from aimbat.logger import logger
    from rich.console import Console

    short = table_parameters.short

    with Session(engine) as session:
        logger.info("Printing data sources table.")

        if all_events:
            aimbat_data_sources = session.exec(select(AimbatDataSource)).all()
            title = "Data sources for all events"
        else:
            event = resolve_event(session, global_parameters.event_id)
            aimbat_data_sources = get_data_for_event(session, event)
            time = event.time.strftime("%Y-%m-%d %H:%M:%S") if short else event.time
            id = uuid_shortener(session, event) if short else event.id
            title = f"Data sources for event {time} (ID={id})"

        logger.debug(f"Found {len(aimbat_data_sources)} data sources in total.")

        rows = [
            [
                uuid_shortener(session, a) if short else str(a.id),
                str(a.datatype),
                str(a.sourcename),
                (
                    uuid_shortener(session, a.seismogram)
                    if short
                    else str(a.seismogram.id)
                ),
            ]
            for a in aimbat_data_sources
        ]

        table = make_table(title=title)

        table.add_column(
            "ID (shortened)" if short else "ID",
            justify="center",
            style=TABLE_STYLING.id,
            no_wrap=True,
        )
        table.add_column("Datatype", justify="center", style=TABLE_STYLING.mine)
        table.add_column(
            "Source", justify="left", style=TABLE_STYLING.mine, no_wrap=True
        )
        table.add_column(
            "Seismogram ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True
        )

        for row in rows:
            table.add_row(*row)

        console = Console()
        console.print(table)

@app.command(name="default")
@simple_exception
def cli_event_default(
    new_default_event_id: Annotated[uuid.UUID, id_parameter(AimbatEvent)],
    *,
    global_parameters: DebugTrait = DebugTrait(),
) -> None:
    """Select the event to be default for processing."""
    from aimbat.core import set_default_event_by_id
    from aimbat.db import engine

    with Session(engine) as session:
        set_default_event_by_id(session, new_default_event_id)

@app.command(name="delete")
@simple_exception
def cli_event_delete(
    event_id: Annotated[uuid.UUID, id_parameter(AimbatEvent)],
    *,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Delete existing event."""
    from aimbat.db import engine
    from aimbat.core import delete_event_by_id

    with Session(engine) as session:
        delete_event_by_id(session, event_id)

@app.command(name="dump")
@simple_exception
def cli_event_dump(
    *,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Dump the contents of the AIMBAT event table to JSON.

    Output can be piped or redirected for use in external tools or scripts.
    """
    from aimbat.db import engine
    from aimbat.core import dump_event_table_to_json
    from rich import print_json

    with Session(engine) as session:
        print_json(dump_event_table_to_json(session))

@app.command(name="list")
@simple_exception
def cli_event_list(
    *,
    table_parameters: TableParameters = TableParameters(),
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Print a table of events stored in the AIMBAT project.

    The default event is highlighted. Use `event default` to change which event
    is processed by subsequent commands.
    """
    from aimbat.db import engine
    from aimbat.core import dump_event_table_to_json
    from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING
    from aimbat.logger import logger
    from pandas import Timestamp

    short = table_parameters.short

    with Session(engine) as session:
        logger.info("Printing AIMBAT events table.")

        json_to_table(
            data=dump_event_table_to_json(session, as_string=False),
            title="AIMBAT Events",
            column_order=[
                "id",
                "is_default",
                "time",
                "latitude",
                "longitude",
                "depth",
                "completed",
                "seismogram_count",
                "station_count",
            ],
            formatters={
                "id": lambda x: (
                    uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x
                ),
                "is_default": TABLE_STYLING.bool_formatter,
                "time": lambda x: TABLE_STYLING.timestamp_formatter(
                    Timestamp(x), short
                ),
                "latitude": lambda x: f"{x:.3f}" if short else str(x),
                "longitude": lambda x: f"{x:.3f}" if short else str(x),
                "depth": lambda x: f"{x:.0f}" if short and x is not None else str(x),
                "completed": TABLE_STYLING.bool_formatter,
                "last_modified": lambda x: TABLE_STYLING.timestamp_formatter(
                    Timestamp(x), short
                ),
            },
            common_column_kwargs={"justify": "center"},
            column_kwargs={
                "id": {
                    "header": "ID (shortened)" if short else "ID",
                    "style": TABLE_STYLING.id,
                    "no_wrap": True,
                },
                "is_default": {
                    "header": "Default",
                    "style": TABLE_STYLING.mine,
                    "no_wrap": True,
                },
                "time": {
                    "header": "Date & Time",
                    "style": TABLE_STYLING.mine,
                    "no_wrap": True,
                },
                "last_modified": {
                    "header": "Last Modified",
                    "style": TABLE_STYLING.mine,
                    "no_wrap": True,
                },
                "latitude": {"style": TABLE_STYLING.mine},
                "longitude": {"style": TABLE_STYLING.mine},
                "depth": {"style": TABLE_STYLING.mine},
                "completed": {"style": TABLE_STYLING.parameters},
                "seismogram_count": {
                    "header": "# Seismograms",
                    "style": TABLE_STYLING.linked,
                },
                "station_count": {
                    "header": "# Stations",
                    "style": TABLE_STYLING.linked,
                },
            },
        )

@parameter.command(name="dump")
@simple_exception
def cli_event_parameter_dump(
    all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Dump event parameter table to json."""
    from aimbat.db import engine
    from aimbat.core import dump_event_parameter_table_to_json, resolve_event
    from sqlmodel import Session
    from rich import print_json

    with Session(engine) as session:
        event = (
            resolve_event(session, global_parameters.event_id)
            if not all_events
            else None
        )
        print_json(
            dump_event_parameter_table_to_json(
                session, all_events, as_string=True, event=event
            )
        )

@parameter.command(name="get")
@simple_exception
def cli_event_parameter_get(
    name: EventParameter,
    *,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Get parameter value for an event.

    Args:
        name: Event parameter name.
    """

    from aimbat.db import engine
    from aimbat.core import get_event_parameter, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        value = get_event_parameter(session, event, name)
        if isinstance(value, Timedelta):
            print(f"{value.total_seconds()}s")
        else:
            print(value)

@parameter.command(name="list")
@simple_exception
def cli_event_parameter_list(
    all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False,
    global_parameters: GlobalParameters = GlobalParameters(),
    table_parameters: TableParameters = TableParameters(),
) -> None:
    """List processing parameter values for the default event.

    Displays all event-level parameters (e.g. time window, bandpass filter
    settings, minimum ccnorm) in a table.
    """

    from aimbat.db import engine
    from aimbat.core import dump_event_parameter_table_to_json, resolve_event
    from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING
    from aimbat.logger import logger

    short = table_parameters.short

    with Session(engine) as session:
        if all_events:
            logger.info("Printing AIMBAT event parameters table for all events.")
            json_to_table(
                data=dump_event_parameter_table_to_json(
                    session, all_events=True, as_string=False
                ),
                title="Event parameters for all events",
                skip_keys=["id"],
                column_order=[
                    "event_id",
                    "completed",
                    "window_pre",
                    "window_post",
                    "min_ccnorm",
                ],
                formatters={
                    "event_id": lambda x: (
                        uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x
                    ),
                },
                common_column_kwargs={"highlight": True},
                column_kwargs={
                    "event_id": {
                        "header": "Event ID (shortened)" if short else "Event ID",
                        "justify": "center",
                        "style": TABLE_STYLING.mine,
                    },
                },
            )
        else:
            logger.info("Printing AIMBAT event parameters table for default event.")

            event = resolve_event(session, global_parameters.event_id)
            json_to_table(
                data=event.parameters.model_dump(mode="json"),
                title=f"Event parameters for event: {uuid_shortener(session, event) if short else str(event.id)}",
                skip_keys=["id", "event_id"],
                common_column_kwargs={"highlight": True},
                column_kwargs={
                    "Key": {
                        "header": "Parameter",
                        "justify": "left",
                        "style": TABLE_STYLING.id,
                    },
                },
            )

@parameter.command(name="set")
@simple_exception
def cli_event_parameter_set(
    name: EventParameter,
    value: str,
    *,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Set parameter value for an event.

    Args:
        name: Event parameter name.
        value: Event parameter value. For timedelta parameters, bare numbers
            are interpreted as seconds.
    """
    from aimbat.db import engine
    from aimbat.core import set_event_parameter, resolve_event
    from sqlmodel import Session

    _TIMEDELTA_PARAMS = (EventParameter.WINDOW_PRE, EventParameter.WINDOW_POST)
    parsed_value: Timedelta | str = value
    if name in _TIMEDELTA_PARAMS:
        try:
            parsed_value = Timedelta(seconds=float(value))
        except ValueError:
            parsed_value = Timedelta(value)

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        set_event_parameter(session, event, name, parsed_value)

@app.command(name="ccnorm")
@simple_exception
def cli_pick_min_ccnorm(
    *,
    iccs_parameters: IccsPlotParameters = IccsPlotParameters(),
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Interactively pick a new minimum cross-correlation norm for auto-selection.

    Opens an interactive plot; click to set the ccnorm threshold. Seismograms
    whose cross-correlation with the stack falls below this value will be
    automatically de-selected when running ICCS with `--autoselect`.
    """
    from aimbat.db import engine
    from aimbat.core import create_iccs_instance, update_min_ccnorm, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        iccs = create_iccs_instance(session, event).iccs
        update_min_ccnorm(
            session,
            event,
            iccs,
            iccs_parameters.context,
            iccs_parameters.all,
            return_fig=False,
        )

@app.command(name="window")
@simple_exception
def cli_pick_timewindow(
    *,
    iccs_parameters: IccsPlotParameters = IccsPlotParameters(),
    use_seismogram_image: Annotated[bool, Parameter(name="img")] = False,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Interactively pick a new cross-correlation time window for an event.

    Opens an interactive plot; click to set the left and right window boundaries,
    then close the window to save. The window controls which portion of the
    seismogram is used during ICCS alignment.

    Args:
        use_seismogram_image: Use the seismogram image to pick the time window.
    """
    from aimbat.db import engine
    from aimbat.core import create_iccs_instance, update_timewindow, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        iccs = create_iccs_instance(session, event).iccs
        update_timewindow(
            session,
            event,
            iccs,
            iccs_parameters.context,
            iccs_parameters.all,
            use_seismogram_image,
            return_fig=False,
        )

@app.command(name="phase")
@simple_exception
def cli_update_phase_pick(
    *,
    iccs_parameters: IccsPlotParameters = IccsPlotParameters(),
    use_seismogram_image: Annotated[bool, Parameter(name="img")] = False,
    global_parameters: GlobalParameters = GlobalParameters(),
) -> None:
    """Interactively pick a new phase arrival time (t1) for the default event.

    Opens an interactive plot; click on the waveform to place the new pick,
    then close the window to save. The pick is stored as `t1` for each
    selected seismogram.

    Args:
        use_seismogram_image: Use the seismogram image to update pick.
    """
    from aimbat.db import engine
    from aimbat.core import create_iccs_instance, update_pick, resolve_event
    from sqlmodel import Session

    with Session(engine) as session:
        event = resolve_event(session, global_parameters.event_id)
        iccs = create_iccs_instance(session, event).iccs
        update_pick(
            session,
            iccs,
            iccs_parameters.context,
            iccs_parameters.all,
            use_seismogram_image,
            return_fig=False,
        )