#!/usr/bin/env python3
"""Run one session digest compression job.

Dry-run is the default: it loads one ended session, releases the DB connection,
calls the configured engine, applies the local digest quality gate, and prints a
metadata-only result. Use --apply to replace the fallback digest after acceptance.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[1]
API_ROOT = REPO_ROOT / "apps" / "api"

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8")


def _prepare_imports() -> None:
    os.chdir(API_ROOT)
    api_root = str(API_ROOT)
    if api_root not in sys.path:
        sys.path.insert(0, api_root)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Run the M2 session digest worker for one session. Dry-run is the "
            "default; --apply writes accepted LLM digest output back to DB."
        )
    )
    parser.add_argument("--session-id", required=True, help="app.sessions.id to compress")
    parser.add_argument("--model", help="optional engine model override")
    parser.add_argument(
        "--forbidden-substring",
        action="append",
        default=[],
        help="raw substring that must reject the candidate digest; repeatable",
    )
    parser.add_argument("--apply", action="store_true", help="write accepted digest output to DB")
    parser.add_argument("--json", action="store_true", help="emit machine-readable JSON")
    parser.add_argument(
        "--show-digest",
        action="store_true",
        help="include the accepted digest text in output; default is metadata only",
    )
    parser.add_argument(
        "--allow-rejected",
        action="store_true",
        help="exit 0 when the engine returned a rejected candidate",
    )
    return parser


async def _run(args: argparse.Namespace) -> dict[str, Any]:
    _prepare_imports()
    from app.db import acquire, close_pool, init_pool
    from app.engine_client import engine_client
    from app.routes.admin import apply_engine_config_from_store
    from app.session_persistence import record_llm_call_audit
    from app.services import session_digest_worker

    started = time.perf_counter()
    await init_pool()
    try:
        async with acquire(role="admin") as conn:
            loaded = await session_digest_worker.load_session_digest_job(conn, args.session_id)
        if loaded is None:
            return {
                "ok": False,
                "found": False,
                "applied": False,
                "session_id": args.session_id,
                "error": "session digest job not found",
            }

        await apply_engine_config_from_store()
        await engine_client.startup()
        worker = await session_digest_worker.run_session_digest_worker(
            loaded.job,
            engine_client,
            existing_case_digest=loaded.existing_case_digest,
            forbidden_substrings=tuple(args.forbidden_substring or ()),
            model=args.model,
            audit_hook=record_llm_call_audit,
        )

        applied = False
        if args.apply and worker.apply_plan is not None:
            async with acquire(role="admin") as conn:
                applied = await session_digest_worker.apply_session_digest_plan(
                    conn,
                    worker.apply_plan,
                    learner_id=loaded.learner_id,
                )
        return _result_payload(
            loaded=loaded,
            worker=worker,
            applied=applied,
            apply_requested=bool(args.apply),
            show_digest=bool(args.show_digest),
            elapsed_ms=int((time.perf_counter() - started) * 1000),
        )
    finally:
        await engine_client.shutdown()
        await close_pool()


def _result_payload(
    *,
    loaded: Any,
    worker: Any,
    applied: bool,
    apply_requested: bool,
    show_digest: bool,
    elapsed_ms: int,
) -> dict[str, Any]:
    quality = worker.outcome.quality
    result = worker.outcome.result
    apply_plan = worker.apply_plan
    payload: dict[str, Any] = {
        "ok": bool(result is not None),
        "found": True,
        "accepted": bool(result is not None),
        "applied": bool(applied),
        "apply_requested": apply_requested,
        "fallback_required": bool(worker.fallback_required),
        "session_id": loaded.job.session_id,
        "case_id": loaded.job.case_id,
        "session_no": loaded.job.session_no,
        "quality": {
            "accepted": bool(quality.accepted),
            "reason": quality.reason,
            "retryable": bool(quality.retryable),
            "details": list(quality.details),
        },
        "engine": {
            "provider": worker.response.provider,
            "model": worker.response.model,
            "tokens_in": worker.response.tokens_in,
            "tokens_out": worker.response.tokens_out,
            "cost_usd": worker.response.cost_usd,
            "inference_geo": worker.response.inference_geo,
        },
        "elapsed_ms": elapsed_ms,
    }
    if apply_plan is not None:
        payload["apply_plan"] = {
            "compressed_by": apply_plan.compressed_by,
            "token_count": apply_plan.token_count,
            "open_thread_count": len(apply_plan.result.open_threads),
            "digest_chars": len(apply_plan.result.digest),
            "case_digest_chars": len(apply_plan.case_digest or ""),
        }
        if show_digest:
            payload["digest"] = apply_plan.result.digest
    return payload


def _print_text(payload: dict[str, Any]) -> None:
    if not payload.get("found"):
        print(payload.get("error") or "session digest job not found", file=sys.stderr)
        return
    status = "accepted" if payload.get("accepted") else "rejected"
    mode = "applied" if payload.get("applied") else "dry-run"
    print(
        "session digest worker: {status} ({mode}) session={session} reason={reason}".format(
            status=status,
            mode=mode,
            session=payload.get("session_id"),
            reason=(payload.get("quality") or {}).get("reason"),
        )
    )
    plan = payload.get("apply_plan") or {}
    if plan:
        print(
            "candidate: compressed_by={compressed_by} token_count={token_count} "
            "digest_chars={digest_chars} case_digest_chars={case_digest_chars}".format(**plan)
        )
    engine = payload.get("engine") or {}
    if engine:
        print(
            "engine: provider={provider} model={model} tokens={tokens_in}/{tokens_out} cost={cost_usd}".format(
                **engine
            )
        )
    if payload.get("digest"):
        print("digest:")
        print(payload["digest"])


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv if argv is not None else sys.argv[1:]))
    try:
        payload = asyncio.run(_run(args))
    except Exception as exc:  # noqa: BLE001 - CLI should report compactly.
        payload = {
            "ok": False,
            "found": False,
            "applied": False,
            "error_type": exc.__class__.__name__,
            "error": str(exc),
        }
    if args.json:
        print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True))
    else:
        _print_text(payload)
    if payload.get("ok") or (args.allow_rejected and payload.get("found")):
        return 0
    if payload.get("found"):
        return 2
    return 3


if __name__ == "__main__":
    raise SystemExit(main())