from __future__ import annotations

import atexit
import json
import uuid
from dataclasses import asdict
from pathlib import Path
from typing import Any, Dict, Optional

from .config_loader import dump_effective_config, load_file_config, normalize_mode
from .llm_middleware import LLMMiddleware, MiddlewareConfig, create_middleware
from .policy import Policy, ProtectMode, apply_overrides, parse_mode
from .runtime_events import JsonlEventSink


def _safe_slug(s: str) -> str:
    out = []
    for ch in s or "":
        if ch.isalnum() or ch in "._-":
            out.append(ch)
        else:
            out.append("_")
    return "".join(out)[:80] or "run"


class RuntimeSession:
    """
    Production-style runtime wrapper:
    - config loading (pyproject/yaml/json) + code overrides
    - policy mode -> middleware config
    - structured events JSONL
    - per-run HTML report (experimental)
    """

    def __init__(
        self,
        *,
        mode: ProtectMode,
        run_dir: Path,
        provider: str = "",
        model: str = "",
        middleware_config: MiddlewareConfig,
        middleware: LLMMiddleware,
        events_path: Path,
        effective_config: Dict[str, Any],
        auto_report: bool = True,
    ):
        self.mode = mode
        self.run_dir = run_dir
        self.provider = provider
        self.model = model
        self.middleware_config = middleware_config
        self.middleware = middleware
        self.events_path = events_path
        self.effective_config = effective_config
        self.auto_report = bool(auto_report)
        self._closed = False

    def __enter__(self) -> "RuntimeSession":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def wrap(self, client: Any, *, provider: Optional[str] = None) -> Any:
        from .llm_middleware import wrap_anthropic, wrap_openai, wrap_openrouter

        prov = (provider or self.provider or "").lower().strip()
        if prov == "openrouter":
            return wrap_openrouter(client, middleware=self.middleware)
        if prov == "openai":
            return wrap_openai(client, middleware=self.middleware)
        if prov in {"anthropic", "claude"}:
            return wrap_anthropic(client, middleware=self.middleware)

        # Best-effort auto detect.
        if hasattr(client, "chat") and getattr(getattr(client, "chat", None), "completions", None):
            return wrap_openai(client, middleware=self.middleware)
        if hasattr(client, "messages"):
            return wrap_anthropic(client, middleware=self.middleware)
        raise TypeError(
            "Unsupported client shape for RuntimeSession.wrap(). Pass provider explicitly."
        )

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if not self.auto_report:
            return
        try:
            from entropic_core.experimental.reporting.runtime_report import generate_runtime_report

            run_meta = {
                "run_id": str(self.middleware_config.run_id or ""),
                "provider": self.provider,
                "model": self.model,
                "mode": self.mode.value,
                "events_path": str(self.events_path),
            }
            generate_runtime_report(
                events_path=self.events_path, outdir=self.run_dir, run_meta=run_meta
            )
        except Exception:
            return


def start_runtime_session(
    *,
    mode: Optional[str],
    run_dir: str = "entropic_runs",
    provider: str = "",
    model: str = "",
    config_path: Optional[str] = None,
    overrides: Optional[Dict[str, Any]] = None,
    auto_report: bool = True,
) -> RuntimeSession:
    file_cfg = load_file_config(config_path)

    # mode: code override > file > default
    m = parse_mode(mode) if mode is not None else normalize_mode(file_cfg)
    policy = Policy(mode=m)

    # Base from policy, then apply file overrides, then apply code overrides.
    cfg = policy.to_middleware_config()
    cfg = apply_overrides(cfg, {k: v for k, v in file_cfg.items() if k in asdict(cfg)})
    if overrides:
        cfg = apply_overrides(cfg, overrides)

    # Run dir structure: <run_dir>/<run_id>_<provider>_<model>/
    rid = uuid.uuid4().hex[:12]
    run_name = _safe_slug(f"{rid}_{provider}_{model}")
    root = Path(run_dir) / run_name
    root.mkdir(parents=True, exist_ok=True)

    events_path = root / "events.jsonl"
    sink = JsonlEventSink(events_path)

    cfg.event_sink = sink
    cfg.run_id = run_name

    # Persist effective config for reproducibility.
    eff = dump_effective_config(
        mode=m, middleware_config=cfg, run_dir=root, extra={"provider": provider, "model": model}
    )
    (root / "effective_config.json").write_text(json.dumps(eff, indent=2), encoding="utf-8")

    mw = create_middleware(cfg)

    sess = RuntimeSession(
        mode=m,
        run_dir=root,
        provider=provider,
        model=model,
        middleware_config=cfg,
        middleware=mw,
        events_path=events_path,
        effective_config=eff,
        auto_report=auto_report,
    )
    return sess


def start_runtime_session_atexit(**kwargs: Any) -> RuntimeSession:
    sess = start_runtime_session(**kwargs)
    atexit.register(sess.close)
    return sess
