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 event_parameter(help: str | None = None) -> Parameter:
    """Return a cyclopts `Parameter` for selecting a single event by UUID or prefix.

    Args:
        help: Custom help string; falls back to a generic event UUID prompt.
    """
    from aimbat.models import AimbatEvent

    return Parameter(
        name=["event", "event-id"],
        help=help or "UUID (or unique prefix) of event to process.",
        env_var="DEFAULT_EVENT_ID",
        converter=_make_uuid_converter(AimbatEvent),
    )

def event_parameter_is_all(event_id: UUID | Literal["all"]) -> TypeIs[Literal["all"]]:
    """Return `True` if `event_id` is the literal string `"all"` (case-insensitive)."""
    if isinstance(event_id, str) and event_id.lower() == "all":
        return True
    return False

def event_parameter_with_all(help: str | None = None) -> Parameter:
    """Return a cyclopts `Parameter` for selecting an event or the literal `"all"`.

    Args:
        help: Custom help string; falls back to a generic prompt.
    """
    from aimbat.models import AimbatEvent

    return Parameter(
        name=["event", "event-id"],
        help=help
        or '"all" for all events, or UUID (or unique prefix) of event to process.',
        env_var="DEFAULT_EVENT_ID",
        converter=_make_uuid_converter(AimbatEvent, allow_all=True),
        show_choices=False,
    )

def id_parameter(model_class: type, help: str = "") -> Parameter:
    """Return a cyclopts `Parameter` for selecting a record by UUID or unique prefix.

    Args:
        model_class: AIMBAT model class used to resolve short UUID prefixes.
        help: Custom help string; falls back to a generic UUID prompt if empty.
    """
    return Parameter(
        name="id",
        help=help or "UUID (or any unique prefix).",
        converter=_make_uuid_converter(model_class),
    )

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 open_in_editor(initial_content: str) -> str:
    """Write `initial_content` to a temporary Markdown file, open it in `$EDITOR`,
    and return the (possibly updated) content after the editor exits.

    The temporary file uses `delete=False` so that it can be opened by a
    second process on Windows (which prohibits opening a file that is already
    open). It is always removed in a `finally` block.

    The editor command is taken from `$EDITOR` or `$VISUAL`. If neither is
    set, `notepad` is used on Windows and `vi` elsewhere. To use a GUI editor
    that does not block by default (e.g. VS Code), set
    ``EDITOR="code --wait"``.
    """
    import os
    import shlex
    import subprocess
    import tempfile

    editor = os.environ.get("EDITOR") or os.environ.get("VISUAL")
    if not editor:
        editor = "notepad" if sys.platform == "win32" else "vi"

    with tempfile.NamedTemporaryFile(
        mode="w",
        suffix=".md",
        delete=False,
        encoding="utf-8",
    ) as tmp:
        tmp.write(initial_content)
        tmp_path = tmp.name

    try:
        subprocess.run([*shlex.split(editor), tmp_path], check=False)
        with open(tmp_path, encoding="utf-8") as f:
            return f.read()
    finally:
        os.unlink(tmp_path)

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

def station_parameter_is_all(
    station_id: UUID | Literal["all"],
) -> TypeIs[Literal["all"]]:
    """Return `True` if `station_id` is the literal string `"all"` (case-insensitive)."""
    if isinstance(station_id, str) and station_id.lower() == "all":
        return True
    return False

def station_parameter_with_all(help: str | None = None) -> Parameter:
    """Return a cyclopts `Parameter` for selecting a station or the literal `"all"`.

    Args:
        help: Custom help string; falls back to a generic prompt.
    """
    from aimbat.models import AimbatStation

    return Parameter(
        name=["station", "station-id"],
        help=help
        or '"all" for all stations, or UUID (or unique prefix) of station to process.',
        converter=_make_uuid_converter(AimbatStation, allow_all=True),
        show_choices=False,
    )

def use_event_parameter() -> Parameter:
    """Return a cyclopts `Parameter` for linking data to an existing event record."""
    from aimbat.models import AimbatEvent

    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(AimbatEvent),
    )

def use_matrix_image() -> Parameter:
    """Return a cyclopts `Parameter` for switching from stack to matrix image plots."""
    return Parameter(
        name="matrix",
        help="Use matrix image instead of stack plot.",
    )

def use_station_parameter() -> Parameter:
    """Return a cyclopts `Parameter` for linking data to an existing station record."""
    from aimbat.models import AimbatStation

    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(AimbatStation),
    )

@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,
        )