"""`aia` — AI Atlas operations CLI."""
from __future__ import annotations

import asyncio
import json
import logging
from pathlib import Path
from typing import Annotated

import typer
from rich.console import Console
from rich.table import Table

from aiatlas.config import settings
from aiatlas.logging import setup_logging

app = typer.Typer(name="aia", help="AI Atlas — the global intelligence layer for AI.", no_args_is_help=True, add_completion=False)
console = Console(stderr=True)
out = Console()


@app.callback()
def _main(verbose: Annotated[bool, typer.Option("--verbose", "-v")] = False) -> None:
    setup_logging(level=logging.DEBUG if verbose else logging.INFO, service="aia-cli")
    settings.ensure_dirs()


def _run(coro):  # type: ignore[no-untyped-def]
    from aiatlas.db import dispose
    from aiatlas.services import cache

    async def wrapper():  # type: ignore[no-untyped-def]
        try:
            return await coro
        finally:
            await dispose()
            await cache.close()

    return asyncio.run(wrapper())


@app.command()
def migrate(revision: str = "head") -> None:
    """Apply database migrations (forward-only)."""
    from alembic import command
    from alembic.config import Config

    root = Path(__file__).resolve().parents[2]
    cfg = Config(str(root / "alembic.ini"))
    cfg.set_main_option("script_location", str(root / "migrations"))
    command.upgrade(cfg, revision)
    out.print("[green]migrations applied[/]")


@app.command()
def seed() -> None:
    """Seed sources, connectors and curated registries (idempotent)."""
    from aiatlas.db import transaction
    from aiatlas.registry.seed import seed as _seed

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await _seed(conn)

    res = _run(go())
    out.print(f"[green]seeded[/] {json.dumps(res)}")


@app.command()
def connectors() -> None:
    """List registered connectors (code) and their database state."""
    from aiatlas.connectors import registry
    from aiatlas.db import fetch_all, transaction

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await fetch_all(conn, "select * from connectors order by priority, name")

    rows = {r["name"]: r for r in _run(go())}
    t = Table(title="connectors")
    for col in ("name", "tier", "enabled", "health", "interval", "last success", "last change", "failures", "parser"):
        t.add_column(col)
    for name, cls in registry().items():
        r = rows.get(name)
        t.add_row(name, str(cls.tier), "✓" if (r and r["enabled"]) else "✗", r["health"] if r else "—", f"{(r['interval_seconds'] if r else cls.interval_seconds) // 60} min",
                  r["last_success_at"].strftime("%m-%d %H:%M") if r and r["last_success_at"] else "—",
                  r["last_change_at"].strftime("%m-%d %H:%M") if r and r["last_change_at"] else "—",
                  str(r["consecutive_failures"]) if r else "—", cls.parser_version)
    out.print(t)


@app.command()
def run(name: str, force: bool = typer.Option(False, "--force", help="ignore enabled/circuit/conditional headers"),
        file: list[str] = typer.Option(None, "--file", help="key=path or url=path local override (fixtures / seed snapshots)"),
        max_targets: int = typer.Option(0, "--max-targets"), url: list[str] = typer.Option(None, "--url", help="only these URLs")) -> None:
    """Run one connector now."""
    from aiatlas.connectors import get

    overrides = dict(f.split("=", 1) for f in (file or []))
    ctx = _run(get(name).run(force=force, file_overrides=overrides, max_targets=max_targets or None, only_urls=url or None))
    out.print(json.dumps({k: v for k, v in ctx.stats.__dict__.items()}, default=str, indent=1))


@app.command()
def reprocess(name: str, url: list[str] = typer.Option(None, "--url")) -> None:
    """Re-extract from stored snapshots (no network) — after a parser improvement."""
    from aiatlas.connectors import get

    ctx = _run(get(name).run(reprocess=True, force=True, only_urls=url or None))
    out.print(json.dumps({k: v for k, v in ctx.stats.__dict__.items()}, default=str, indent=1))


@app.command()
def crawl(priority: int = typer.Option(9, "--priority", help="run connectors with priority <= this"), force: bool = False,
          only: list[str] = typer.Option(None, "--only")) -> None:
    """Run every enabled connector once, in priority order (initial corpus build)."""
    from aiatlas.connectors import get, registry
    from aiatlas.db import fetch_all, transaction

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            rows = await fetch_all(conn, "select name, enabled, priority from connectors order by priority, name")
        results = {}
        for r in rows:
            if r["name"] not in registry() or (not r["enabled"] and not force) or r["priority"] > priority or (only and r["name"] not in only):
                continue
            console.print(f"[cyan]▶ {r['name']}[/]")
            try:
                ctx = await get(r["name"]).run(force=force)
                results[r["name"]] = {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}
            except Exception as exc:  # noqa: BLE001
                results[r["name"]] = {"error": str(exc)[:300]}
        return results

    out.print(json.dumps(_run(go()), default=str, indent=1))


@app.command()
def status() -> None:
    """Connector health table and queue depth."""
    from aiatlas.db import fetch_all, transaction
    from aiatlas.services.jobs import queue_depth

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            rows = await fetch_all(conn, """select c.name, c.enabled, c.health, c.last_success_at, c.next_run_at, c.consecutive_failures,
                                            (select count(*) from documents d where d.connector_name = c.name) as docs,
                                            (select count(*) from snapshots s join documents d on d.id = s.document_id where d.connector_name = c.name) as snaps,
                                            (select status from connector_runs r where r.connector_name = c.name order by started_at desc limit 1) as last_status
                                            from connectors c order by c.priority, c.name""")
            return rows, await queue_depth(conn)

    rows, depth = _run(go())
    t = Table(title="AI Atlas connectors")
    for col in ("connector", "on", "health", "last run", "last success", "next run", "docs", "snapshots", "fails"):
        t.add_column(col)
    for r in rows:
        t.add_row(r["name"], "✓" if r["enabled"] else "✗", r["health"], r["last_status"] or "—",
                  r["last_success_at"].strftime("%m-%d %H:%M") if r["last_success_at"] else "—",
                  r["next_run_at"].strftime("%m-%d %H:%M") if r["next_run_at"] else "—", str(r["docs"]), str(r["snaps"]), str(r["consecutive_failures"]))
    out.print(t)
    out.print(f"queue: {json.dumps(depth)}")


@app.command()
def stats() -> None:
    """Live counters (from the database) and archive size."""
    from aiatlas.services.stats import compute_stats

    out.print(json.dumps(_run(compute_stats()), default=str, indent=1))


@app.command()
def quality(limit: int = 20000) -> None:
    """Recompute entity quality scores."""
    from aiatlas.services.quality import recompute

    out.print(json.dumps(_run(recompute(limit=limit))))


@app.command()
def schedule(no_worker: bool = typer.Option(False, "--no-worker")) -> None:
    """Long-running scheduler: due connectors, job worker, hourly stats/quality, nightly backup."""
    from aiatlas.services.scheduler import main

    setup_logging(service="aia-scheduler")
    asyncio.run(main(with_worker=not no_worker))


@app.command()
def worker(concurrency: int = typer.Option(0, "--concurrency"), kind: list[str] = typer.Option(None, "--kind")) -> None:
    """Job worker only (LLM extraction, embeddings, reprocessing). Run on any node with database access."""
    from aiatlas.services.jobs import run_worker

    setup_logging(service="aia-worker")
    asyncio.run(run_worker(concurrency=concurrency or None, kinds=kind or None))


@app.command()
def api(host: str = "", port: int = 0, reload: bool = False, workers: int = 1) -> None:
    """Serve the FastAPI application."""
    import uvicorn

    uvicorn.run("aiatlas.api.main:app", host=host or settings.api_host, port=port or settings.api_port, reload=reload, workers=workers if not reload else 1,
                log_level="info", access_log=False, proxy_headers=True)


@app.command()
def backup() -> None:
    """pg_dump into AIA_DATA_DIR/backups (keeps 30)."""
    from aiatlas.services.backup import backup_database

    out.print(str(backup_database()))


@app.command()
def llm(text: str = typer.Argument("", help="optional text to classify"), health: bool = typer.Option(True)) -> None:
    """Check the local LLM factory gateway (and optionally classify a snippet)."""
    from aiatlas.services.llm import gateway

    async def go():  # type: ignore[no-untyped-def]
        h = await gateway.health()
        if text:
            h["classification"] = await gateway.classify(text=text, labels=["model_release", "pricing", "research_paper", "company_news", "other"])
        return h

    out.print(json.dumps(_run(go()), indent=1))


@app.command()
def embed(limit: int = 200) -> None:
    """Embed entities missing vectors (requires the LLM gateway)."""
    from aiatlas.services.embeddings import embed_entities, pending_entity_ids

    async def go():  # type: ignore[no-untyped-def]
        ids = await pending_entity_ids(limit)
        return await embed_entities(ids) if ids else {"embedded": 0}

    out.print(json.dumps(_run(go())))


@app.command()
def search(q: str, limit: int = 10) -> None:
    """Search entities (natural-language filters are compiled deterministically)."""
    from aiatlas.db import transaction
    from aiatlas.services.search import compile_query, search_entities

    async def go():  # type: ignore[no-untyped-def]
        query = compile_query(q)
        async with transaction() as conn:
            return query.as_dict(), await search_entities(conn, query, limit=limit)

    query, rows = _run(go())
    out.print(json.dumps(query))
    for r in rows:
        out.print(f"  {r['entity_type']:<10} {r['canonical_name']:<50} {r.get('organization_name') or ''}  /{r['slug']}")


@app.command()
def review(status: str = "pending", limit: int = 30) -> None:
    """Show the human review queue."""
    from aiatlas.db import fetch_all, transaction

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await fetch_all(conn, "select id, kind, reason, created_at from review_queue where status = :s order by created_at desc limit :n", s=status, n=limit)

    for r in _run(go()):
        out.print(f"{r['created_at']:%m-%d %H:%M} {r['kind']:<16} {r['reason']}")


@app.command()
def enqueue_llm(limit: int = 500, task: str = "auto") -> None:
    """Queue LLM extraction for snapshots that are still `llm_pending`/`stored` on documents flagged needs_llm."""
    from aiatlas.db import fetch_all, transaction
    from aiatlas.services.jobs import enqueue

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            rows = await fetch_all(conn, """select s.id, d.entity_id, d.connector_name from snapshots s join documents d on d.id = s.document_id
                                            where s.processing_status in ('llm_pending','stored') and d.needs_llm and s.changed
                                            and s.id = (select id from snapshots s2 where s2.document_id = s.document_id and s2.changed order by observed_at desc limit 1)
                                            order by s.observed_at desc limit :n""", n=limit)
            n = 0
            for r in rows:
                if await enqueue(conn, "llm_extract", {"snapshot_id": r["id"], "task": task, "entity_id": r["entity_id"], "connector": r["connector_name"]},
                                 priority=5, dedupe_key=f"llm_extract:{r['id']}"):
                    n += 1
            return {"queued": n, "candidates": len(rows)}

    out.print(json.dumps(_run(go())))


@app.command()
def canonicalize(apply: bool = typer.Option(False, "--apply", help="write changes (default: dry-run report)"),
                 step: list[str] = typer.Option(None, "--step", help="run only these steps (repeatable); default = all, in canonical order"),
                 as_json: bool = typer.Option(False, "--json", help="machine-readable report")) -> None:
    """Canonicalize the corpus: duplicates → effort variants → artifacts → families → licences → taxonomy → results → events → anomalies.
    Dry-run by default; idempotent under --apply (a second pass reports 0 changes). Never deletes rows, never touches snapshots."""
    from aiatlas.services.canonical import canonicalize as _canon

    report = _run(_canon(apply=apply, steps=step or None))
    out.print(json.dumps(report.as_dict(), indent=1, default=str) if as_json else report.render())


@app.command()
def anomalies(severity: str = typer.Option("", "--severity", help="critical|warning|info"), status: str = typer.Option("open", "--status"),
              limit: int = typer.Option(100, "--limit"), refresh: bool = typer.Option(False, "--refresh", help="re-run the checks before listing")) -> None:
    """List data anomalies (flags with evidence — nothing is deleted or silently fixed)."""
    from aiatlas.db import transaction
    from aiatlas.services.anomalies import list_anomalies, run_checks

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            summary = await run_checks(conn) if refresh else None
            return summary, await list_anomalies(conn, severity=severity or None, status=status, limit=limit)

    summary, rows = _run(go())
    if summary:
        console.print(f"[cyan]checks:[/] {json.dumps(summary)}")
    t = Table(title=f"anomalies ({status}{', ' + severity if severity else ''})")
    for col in ("severity", "check", "entity", "message", "last seen"):
        t.add_column(col)
    for r in rows:
        t.add_row(r["severity"], r["check_name"], r["slug"] or (r["entity_id"] or "—"), r["message"][:110], r["last_seen_at"].strftime("%m-%d %H:%M"))
    out.print(t)


quarantine_app = typer.Typer(help="Runs held by the anomaly detector (nothing written until released).", no_args_is_help=True)
app.add_typer(quarantine_app, name="quarantine")


@quarantine_app.command("list")
def quarantine_list(status: str = typer.Option("pending", "--status"), limit: int = 50) -> None:
    """Show quarantined runs."""
    from aiatlas.db import transaction
    from aiatlas.services.canonical import list_quarantine

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await list_quarantine(conn, status=status, limit=limit)

    t = Table(title=f"quarantined runs ({status or 'all'})")
    for col in ("id", "connector", "created", "status", "docs", "reason"):
        t.add_column(col)
    for r in _run(go()):
        t.add_row(r["id"], r["connector_name"], r["created_at"].strftime("%m-%d %H:%M"), r["status"], str(r["documents"]), r["reason"][:90])
    out.print(t)


@quarantine_app.command("release")
def quarantine_release(quarantine_id: str) -> None:
    """Write the held facts as the connector would have."""
    from aiatlas.db import transaction
    from aiatlas.services.canonical import release_quarantine

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await release_quarantine(conn, quarantine_id, actor="cli")

    out.print(json.dumps(_run(go()), default=str))


@quarantine_app.command("discard")
def quarantine_discard(quarantine_id: str, note: str = typer.Option("", "--note")) -> None:
    """Drop the held facts (the raw snapshots stay archived)."""
    from aiatlas.db import transaction
    from aiatlas.services.canonical import discard_quarantine

    async def go():  # type: ignore[no-untyped-def]
        async with transaction() as conn:
            return await discard_quarantine(conn, quarantine_id, actor="cli", note=note or None)

    out.print(json.dumps(_run(go()), default=str))


if __name__ == "__main__":
    app()