aimbat._cli

@app.command(name="iccs")
@simple_exception
def cli_iccs_run(
    event_id: Annotated[UUID, event_parameter()],
    *,
    autoselect: Annotated[
        bool,
        Parameter(
            name="autoselect",
            help="Whether to automatically de-select seismograms whose"
            " cross-correlation with the stack falls below `min_cc`, and"
            " re-select them if the cross-correlation later exceeds `min_cc`.",
        ),
    ] = False,
    autoflip: Annotated[
        bool,
        Parameter(
            name="autoflip",
            help="Whether to automatically flip seismograms (multiply data"
            " by -1) when the cross-correlation is negative.",
        ),
    ] = False,
    _: DebugParameter = DebugParameter(),
) -> None:
    """Run the ICCS algorithm to align seismograms for an 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.
    """
    from sqlmodel import Session

    from aimbat.core import create_iccs_instance, resolve_event, run_iccs
    from aimbat.db import engine

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

@app.command(name="mccc")
@simple_exception
def cli_mccc_run(
    event_id: Annotated[UUID, event_parameter()],
    *,
    all_seismograms: Annotated[
        bool,
        Parameter(
            name="all",
            help="Include all seismograms of an event in MCCC processing, "
            "not just the currently selected ones.",
        ),
    ] = False,
    _: DebugParameter = DebugParameter(),
) -> None:
    """Run the MCCC algorithm to refine arrival time picks for an event.

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

    from aimbat.core import create_iccs_instance, resolve_event, run_mccc
    from aimbat.db import engine

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

def json_to_table(
    data: dict[str, Any] | list[dict[str, Any]],
    model: type[BaseModel],
    title: str | None = None,
    raw: bool = False,
    col_specs: dict[str, RichColSpec] | None = None,
    column_order: list[str] | None = None,
    key_header: str = "Property",
    value_header: str = "Value",
) -> None:
    """Print a JSON dict or list of dicts as a rich table driven by a Pydantic model.

    Args:
        data: A single row (dict) or list of rows to display.
        model: Pydantic model whose field metadata drives column configuration.
        title: Optional table title.
        raw: If `True`, ignore `RichColSpec` metadata and render using only
            type-based heuristics. Useful for a quick unformatted view.
        col_specs: Optional per-field overrides. Each entry is merged on top of
            the spec derived from the model field, so only the attributes that
            differ need to be set. Ignored when `raw=True`.
        column_order: Optional list of field names that should appear first, in
            that order. Fields not listed appear after in model-declaration order.
        key_header: Header for the property-name column in vertical (dict) tables.
        value_header: Header for the value column in vertical (dict) tables.
    """
    console = Console()
    table = Table(title=title)

    data_list = [data] if isinstance(data, dict) else data
    if not data_list:
        console.print(table)
        return

    # Build specs from model field metadata (skipped in raw mode).
    specs: dict[str, RichColSpec] = {}
    if not raw:
        for name, field in model.model_fields.items():
            spec: RichColSpec | None = None
            # Pydantic Field uses json_schema_extra; SQLModel Field uses schema_extra
            # which spreads its keys directly into _attributes_set on the FieldInfo.
            for source in (
                field.json_schema_extra,
                getattr(field, "_attributes_set", None),
            ):
                if isinstance(source, dict):
                    candidate = source.get("rich")
                    if isinstance(candidate, RichColSpec):
                        spec = candidate
                        break
            specs[name] = spec if spec is not None else RichColSpec()

        # Apply caller-supplied overrides. model_fields_set tracks which fields
        # were explicitly provided, so only those replace the model-derived spec.
        if col_specs:
            for name, override in col_specs.items():
                base = specs.get(name, RichColSpec())
                specs[name] = base.model_copy(
                    update={k: getattr(override, k) for k in override.model_fields_set}
                )

    def _ordered(names: list[str]) -> list[str]:
        if not column_order:
            return names
        ordered = [n for n in column_order if n in names]
        rest = [n for n in names if n not in set(column_order)]
        return ordered + rest

    def _fmt_val(name: str, val: Any) -> str:
        if raw:
            return "" if val is None else str(val)
        if val is None or val is NaT:
            return _MISSING_MARKER
        spec = specs.get(name)
        if spec and spec.formatter:
            return spec.formatter(val)
        if isinstance(val, bool):
            return fmt_bool(val)
        if isinstance(val, float):
            return fmt_float(val)
        if isinstance(val, Timedelta):
            return fmt_timedelta(val)
        if isinstance(val, (Timestamp, datetime)):
            return fmt_timestamp(val)
        low_key = name.lower()
        if (
            isinstance(val, str)
            and val.strip()
            and any(k in low_key for k in ("time", "date", "modified"))
        ):
            try:
                dt = to_datetime(val)
                return fmt_timestamp(dt)
            except (ValueError, TypeError):
                return str(val)
        return str(val)

    field_names = _ordered(list(model.model_fields.keys()))

    if isinstance(data, dict):
        # Vertical table
        table.add_column(key_header, style="cyan")
        table.add_column(value_header)
        for name in field_names:
            if name not in data:
                continue
            spec = specs.get(name)
            if spec and spec.display_title:
                label = spec.display_title
            else:
                field = model.model_fields[name]
                label = field.title if field.title else name
            table.add_row(label, _fmt_val(name, data[name]))
    else:
        # Horizontal table — restrict to fields present in data.
        data_keys: set[str] = set()
        for item in data_list:
            data_keys.update(item.keys())
        visible_fields = [n for n in field_names if n in data_keys]

        for name in visible_fields:
            spec = specs.get(name)
            field = model.model_fields[name]
            header = (
                spec.display_title
                if spec and spec.display_title
                else field.title or name
            )
            kwargs: dict[str, Any] = {}
            default_justify = _justify_for_annotation(field.annotation)
            if default_justify:
                kwargs["justify"] = default_justify
            if spec:
                if spec.justify:
                    kwargs["justify"] = spec.justify
                if spec.style:
                    kwargs["style"] = spec.style
                if spec.no_wrap is not None:
                    kwargs["no_wrap"] = spec.no_wrap
                if spec.highlight is not None:
                    kwargs["highlight"] = spec.highlight
            table.add_column(header, **kwargs)

        for item in data:
            table.add_row(*[_fmt_val(n, item.get(n)) for n in visible_fields])

    console.print(table)

def print_error_panel(e: Exception) -> None:
    """Print an exception to the console in a red panel."""
    from rich.console import Console
    from rich.panel import Panel

    console = Console(stderr=True)
    panel = Panel(
        f"{e}",
        title="Error",
        title_align="left",
        border_style="red",
        expand=True,
    )
    console.print(panel)

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.
    """
    import sys
    from functools import wraps

    @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:
            print_error_panel(e)
            sys.exit(1)

    return wrapper  # type: ignore

@app.command(name="add")
@simple_exception
def cli_data_add(
    data_sources: Annotated[
        list[Path],
        Parameter(
            name="sources",
            help="One or more data source paths to add.",
            consume_multiple=1,
            negative_iterable=(),
            validator=validators.Path(exists=True),
        ),
    ],
    *,
    data_type: Annotated[
        DataType,
        Parameter(
            name="type",
            help="Format of the data sources. Determines which metadata"
            " (station, event, seismogram) can be extracted automatically.",
        ),
    ] = DataType.SAC,
    station_id: Annotated[uuid.UUID | None, use_station_parameter()] = None,
    event_id: Annotated[uuid.UUID | None, use_event_parameter()] = None,
    dry_run: Annotated[
        bool,
        Parameter(
            name="dry-run",
            help="Preview which records would be added without modifying the database.",
        ),
    ] = False,
    show_progress_bar: Annotated[
        bool,
        Parameter(
            name="progress", help="Display a progress bar while ingesting sources."
        ),
    ] = True,
    _: DebugParameter = DebugParameter(),
) -> 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.
    """
    from aimbat.core import add_data_to_project
    from aimbat.db import engine

    disable_progress_bar = not show_progress_bar

    with Session(engine) as session:
        if dry_run:
            results = add_data_to_project(
                session,
                data_sources,
                data_type,
                station_id=station_id,
                event_id=event_id,
                dry_run=True,
                disable_progress_bar=disable_progress_bar,
            )
            _print_dry_run_results(*results)
        else:
            add_data_to_project(
                session,
                data_sources,
                data_type,
                station_id=station_id,
                event_id=event_id,
                dry_run=False,
                disable_progress_bar=disable_progress_bar,
            )

@app.command(name="dump")
@simple_exception
def cli_data_dump(
    *,
    dump_parameters: JsonDumpParameters = JsonDumpParameters(),
) -> None:
    """Dump AIMBAT datasources table as a JSON string.

    Output can be piped or redirected for use in external tools or scripts.
    """
    from rich import print_json

    from aimbat.core import dump_data_table
    from aimbat.db import engine

    with Session(engine) as session:
        print_json(data=dump_data_table(session, by_alias=dump_parameters.by_alias))

@app.command(name="list")
@simple_exception
def cli_data_list(
    event_id: Annotated[uuid.UUID | Literal["all"], event_parameter_with_all()],
    *,
    table_parameters: TableParameters = TableParameters(),
) -> None:
    """Print a table of data sources registered in the AIMBAT project."""
    from aimbat.core import dump_data_table, resolve_event
    from aimbat.db import engine
    from aimbat.logger import logger
    from aimbat.models import AimbatDataSource, AimbatSeismogram, RichColSpec
    from aimbat.utils import uuid_shortener

    from .common import json_to_table

    raw = table_parameters.raw

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

        if event_parameter_is_all(event_id):
            data = dump_data_table(session)
            title = "Data sources for all events"
        else:
            event = resolve_event(session, event_id)
            data = dump_data_table(session, event.id)
            _time = event.time.strftime("%Y-%m-%d %H:%M:%S") if not raw else event.time
            _id = uuid_shortener(session, event) if not raw else event.id
            title = f"Data sources for event {_time} (ID={_id})"

        col_specs = {
            "id": RichColSpec(
                formatter=lambda x: uuid_shortener(
                    session, AimbatDataSource, str_uuid=x
                )
            ),
            "seismogram_id": RichColSpec(
                formatter=lambda x: uuid_shortener(
                    session, AimbatSeismogram, str_uuid=x
                )
            ),
        }

        json_to_table(
            model=AimbatDataSource,
            data=data,
            title=title,
            raw=raw,
            col_specs=col_specs,
        )

@app.command(name="delete")
@simple_exception
def cli_event_delete(
    event_id: Annotated[uuid.UUID, event_parameter()],
    *,
    _: DebugParameter = DebugParameter(),
) -> None:
    """Delete existing event."""
    from aimbat.core import delete_event, resolve_event
    from aimbat.db import engine

    with Session(engine) as session:
        event = resolve_event(session, event_id)
        delete_event(session, event.id)

@app.command(name="dump")
@simple_exception
def cli_event_dump(
    *, dump_parameters: JsonDumpParameters = JsonDumpParameters()
) -> 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 rich import print_json

    from aimbat.core import dump_event_table
    from aimbat.db import engine

    with Session(engine) as session:
        print_json(dump_event_table(session, by_alias=dump_parameters.by_alias))

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

    from aimbat.core import dump_event_table
    from aimbat.db import engine
    from aimbat.logger import logger
    from aimbat.models import AimbatEventRead

    from .common import json_to_table

    if raw := table_parameters.raw:
        exclude = {"short_id"}
    else:
        exclude = {"id"}

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

        json_to_table(
            data=dump_event_table(
                session, from_read_model=True, by_title=False, exclude=exclude
            ),
            model=AimbatEventRead,
            title="AIMBAT Events",
            raw=raw,
        )

@_parameter.command(name="dump")
@simple_exception
def cli_event_parameter_dump(
    *,
    dump_parameters: JsonDumpParameters = JsonDumpParameters(),
) -> None:
    """Dump event parameter table to json."""
    from rich import print_json
    from sqlmodel import Session

    from aimbat.core import dump_event_parameter_table
    from aimbat.db import engine

    by_alias = dump_parameters.by_alias

    with Session(engine) as session:
        print_json(data=dump_event_parameter_table(session, by_alias=by_alias))

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

    Args:
        name: Event parameter name.
    """

    from sqlmodel import Session

    from aimbat.core import resolve_event
    from aimbat.db import engine

    event_id = event_debug_parameters.event_id

    with Session(engine) as session:
        event = resolve_event(session, event_id)
        value = event.parameters.model_dump(mode="json").get(name)
        print(value)

@_parameter.command(name="list")
@simple_exception
def cli_event_parameter_list(
    event_id: Annotated[uuid.UUID | Literal["all"], event_parameter_with_all()],
    *,
    table_parameters: TableParameters = TableParameters(),
) -> None:
    """List processing parameter for an event or all events.

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

    from aimbat.core import dump_event_parameter_table, resolve_event
    from aimbat.db import engine
    from aimbat.models import AimbatEventParameters, RichColSpec
    from aimbat.utils import uuid_shortener

    from .common import json_to_table

    raw = table_parameters.raw

    with Session(engine) as session:
        if event_parameter_is_all(event_id):
            title = "Event parameters for all events"
            data = dump_event_parameter_table(session)
        else:
            event = resolve_event(session, event_id)
            title = f"Event parameters for event: {uuid_shortener(session, event) if not raw else str(event.id)}"
            data = dump_event_parameter_table(
                session, event_id=event.id, exclude={"event_id"}
            )

        column_order = ["id"]
        col_specs = {
            "id": RichColSpec(
                formatter=lambda x: uuid_shortener(
                    session, AimbatEventParameters, str_uuid=x
                )
            ),
            "event_id": RichColSpec(
                formatter=lambda x: uuid_shortener(session, AimbatEvent, str_uuid=x)
            ),
        }

        json_to_table(
            model=AimbatEventParameters,
            title=title,
            data=data,
            raw=raw,
            col_specs=col_specs,
            column_order=column_order,
        )

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

    Args:
        name: Event parameter name.
        value: New parameter value.
    """
    from sqlmodel import Session

    from aimbat.core import resolve_event, set_event_parameter
    from aimbat.db import engine

    event_id = event_debug_parameters.event_id

    with Session(engine) as session:
        event = resolve_event(session, event_id)
        set_event_parameter(session, event.id, name, value, validate_iccs=True)

@app.command(name="cc")
@simple_exception
def cli_pick_min_cc(
    event_id: Annotated[UUID, event_parameter()],
    *,
    iccs_plot_parameters: IccsPlotParameters = IccsPlotParameters(),
    use_matrix_image: Annotated[bool, use_matrix_image()] = True,
    _: DebugParameter = DebugParameter(),
) -> None:
    """Interactively pick a new minimum cross-correlation for auto-selection.

    Opens an interactive plot; click to set the cc threshold. Seismograms
    whose cross-correlation with the stack falls below this value will be
    automatically de-selected when running ICCS with `--autoselect`.
    """
    from sqlmodel import Session

    from aimbat.core import create_iccs_instance, resolve_event
    from aimbat.db import engine
    from aimbat.plot import update_min_cc

    with Session(engine) as session:
        event = resolve_event(session, event_id)
        iccs = create_iccs_instance(session, event).iccs
        update_min_cc(
            session,
            event,
            iccs,
            iccs_plot_parameters.context,
            all_seismograms=iccs_plot_parameters.all_seismograms,
            return_fig=False,
        )

@app.command(name="window")
@simple_exception
def cli_pick_timewindow(
    event_id: Annotated[UUID, event_parameter()],
    *,
    iccs_plot_parameters: IccsPlotParameters = IccsPlotParameters(),
    use_matrix_image: Annotated[bool, use_matrix_image()] = False,
    _: DebugParameter = DebugParameter(),
) -> 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 each
    seismogram is used during ICCS alignment.
    """
    from sqlmodel import Session

    from aimbat.core import create_iccs_instance, resolve_event
    from aimbat.db import engine
    from aimbat.plot import update_timewindow

    with Session(engine) as session:
        event = resolve_event(session, event_id)
        iccs = create_iccs_instance(session, event).iccs
        update_timewindow(
            session,
            event,
            iccs,
            iccs_plot_parameters.context,
            all_seismograms=iccs_plot_parameters.all_seismograms,
            use_matrix_image=use_matrix_image,
            return_fig=False,
        )