"""Trace tools for the tac-qlib-rd MCP server — replace the `trace.sh`/`trace_db.py`/`git_exp.sh` scripts.

The R&D lineage (`/rd/lineage`) and the round book build on the `rd_experiments`
Postgres table. This module exposes the full trace lifecycle as MCP tools so an
agent can drive tracing through the long-lived tac-qlib-rd server instead of
shelling out to bash scripts (which re-import psycopg + reconnect per call and
force the agent to parse prose output).

Because the server is long-lived, `psycopg` is imported and the DB connection
is opened once per call (not once per script invocation), and every tool returns
a single JSON object — no output parsing, fully deterministic.

Git operations (fork / commit / push on the `experiments/` clone) are performed
via `git` subprocess with the repo's mandated credential helper, exactly as the
old `git_exp.sh` did.

Env (from the repo `.env`, already loaded by rd_server): DATABASE_URL,
EMBEDDING_API_BASE_URL, EMBEDDING_API_KEY, GIT_USER, GIT_PASS, GIT_REPO_URL,
TAC_LAKE_DIR.
"""

from __future__ import annotations

import json
import os
import sqlite3
import subprocess
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import psycopg
from psycopg.rows import dict_row

from tac_qlib.trace_embed import embed

EMBEDDING_DIM = 384
_MIN_SCORE = 0.5


# --------------------------------------------------------------------------- db
def _conn():
    url = (os.environ.get("DATABASE_URL") or "").strip()
    if not url:
        raise RuntimeError("DATABASE_URL is not set")
    return psycopg.connect(url, row_factory=dict_row)


def _now_iso() -> str:
    from datetime import datetime, timezone

    return datetime.now(timezone.utc).isoformat()


def _vector_literal(vec: Optional[List[float]]) -> Optional[str]:
    if not vec:
        return None
    return "[" + ",".join(repr(float(v)) for v in vec) + "]"


def _jsonable(obj: Any) -> Any:
    if isinstance(obj, dict):
        return {k: _jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple, set)):
        return [_jsonable(v) for v in obj]
    if hasattr(obj, "isoformat"):
        return obj.isoformat()
    return obj


def _row_json(row: Dict[str, Any]) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    for k, v in row.items():
        if k in ("rational_embedding", "details_embedding"):
            out[k] = v.tolist() if hasattr(v, "tolist") else v
        elif k == "metrics" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:  # noqa: BLE001
                out[k] = v
        else:
            out[k] = v
    return out


def _get_row(exp_id: int) -> Optional[Dict[str, Any]]:
    with _conn() as conn, conn.cursor() as cur:
        cur.execute("SELECT * FROM rd_experiments WHERE id = %s", (exp_id,))
        return cur.fetchone()


def _all_rows(limit: int) -> List[Dict[str, Any]]:
    with _conn() as conn, conn.cursor() as cur:
        cur.execute("SELECT * FROM rd_experiments ORDER BY id DESC LIMIT %s", (limit,))
        return cur.fetchall()


def _init_db() -> None:
    with _conn() as conn, conn.cursor() as cur:
        cur.execute("SELECT 1 FROM pg_extension WHERE extname = 'vector'")
        if not cur.fetchone():
            raise RuntimeError("pgvector extension is not installed. Run: CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute("SELECT to_regclass('public.rd_experiments')")
        exists = bool(cur.fetchone())
        if not exists:
            DDL = """
            CREATE TABLE IF NOT EXISTS rd_experiments (
                id bigserial PRIMARY KEY NOT NULL,
                experiment_name text,
                rational text NOT NULL,
                rational_embedding vector(384),
                details text,
                details_embedding vector(384),
                evaluation text,
                metrics jsonb,
                evolved_from bigint,
                start_ts timestamptz DEFAULT now() NOT NULL,
                end_ts timestamptz,
                git_branch text NOT NULL,
                experiment_ref_id text,
                session_id text,
                mlruns_dir text,
                status text DEFAULT 'starting' NOT NULL,
                created_at timestamptz DEFAULT now() NOT NULL,
                updated_at timestamptz DEFAULT now() NOT NULL
            );
            """
            for stmt in DDL.split(";"):
                stmt = stmt.strip()
                if stmt:
                    cur.execute(stmt)
        # Idempotent backfills so pre-existing tables gain new columns.
        for stmt in ["ALTER TABLE rd_experiments ADD COLUMN IF NOT EXISTS session_id text;"]:
            stmt = stmt.strip()
            if stmt:
                cur.execute(stmt)
        conn.commit()


def _search_evolved_from(text: str, limit: int = 5) -> Optional[int]:
    vec = embed(text)
    if not vec:
        return None
    lit = _vector_literal(vec)
    with _conn() as conn, conn.cursor() as cur:
        cur.execute(
            """
            SELECT id, 1 - LEAST(
                COALESCE(rational_embedding <=> %s::vector, 1),
                COALESCE(details_embedding   <=> %s::vector, 1)
            ) AS similarity
            FROM rd_experiments
            ORDER BY similarity DESC
            LIMIT %s
            """,
            (lit, lit, limit),
        )
        rows = cur.fetchall()
    for row in rows:
        if row["similarity"] is not None and float(row["similarity"]) >= _MIN_SCORE:
            return int(row["id"])
    return None


def _search(query: str, limit: int = 10, min_score: float = _MIN_SCORE) -> List[Dict[str, Any]]:
    vec = embed(query)
    if not vec:
        needle = f"%{query.replace('%', ' ').strip()}%"
        with _conn() as conn, conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, rational, details, git_branch, experiment_ref_id, status,
                       start_ts, end_ts, evaluation, session_id
                FROM rd_experiments
                WHERE rational ILIKE %s OR details ILIKE %s
                ORDER BY id DESC LIMIT %s
                """,
                (needle, needle, limit),
            )
            return [_row_json(r) for r in cur.fetchall()]
    lit = _vector_literal(vec)
    with _conn() as conn, conn.cursor() as cur:
        cur.execute(
            """
            SELECT id, rational, details, git_branch, experiment_ref_id, status,
                   start_ts, end_ts, evaluation, session_id,
                   1 - LEAST(
                       COALESCE(rational_embedding <=> %s::vector, 1),
                       COALESCE(details_embedding   <=> %s::vector, 1)
                   ) AS similarity
            FROM rd_experiments
            ORDER BY similarity DESC
            LIMIT %s
            """,
            (lit, lit, limit),
        )
        rows = cur.fetchall()
    out = []
    for r in rows:
        sim = float(r.get("similarity") or 0)
        if sim < min_score:
            continue
        r = dict(r)
        r["similarity"] = sim
        out.append(_row_json(r))
    return out


def _start(
    rational: str,
    details: str = "",
    evolved_from: str = "none",
    experiment_name: str = "",
    branch: str = "",
    session_id: str = "",
) -> Dict[str, Any]:
    rational = rational.strip()
    details = (details or "").strip()
    if not rational:
        raise ValueError("--rational is required")

    rational_vec = _vector_literal(embed(rational))
    details_vec = _vector_literal(embed(details)) if details else None

    evo: Optional[int] = None
    if evolved_from == "auto":
        evo = _search_evolved_from(f"{rational}\n{details}") if rational_vec or details_vec else None
    elif evolved_from and evolved_from.isdigit():
        evo = int(evolved_from)

    with _conn() as conn, conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO rd_experiments
                (experiment_name, rational, rational_embedding, details, details_embedding,
                 evolved_from, start_ts, git_branch, status, session_id)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'starting', %s)
            RETURNING id
            """,
            (
                experiment_name or None,
                rational,
                rational_vec,
                details,
                details_vec,
                evo,
                _now_iso(),
                branch or "",
                session_id.strip() or None,
            ),
        )
        row = cur.fetchone()
        exp_id = int(row["id"])
        conn.commit()

    branch = branch or f"exp/{exp_id}"
    with _conn() as conn, conn.cursor() as cur:
        cur.execute("UPDATE rd_experiments SET git_branch = %s WHERE id = %s", (branch, exp_id))
        conn.commit()
    return _row_json(_get_row(exp_id) or {})


def _finish(
    exp_id: int,
    ref_id: str = "",
    evaluation: Optional[str] = None,
    metrics: Optional[str] = None,
    mlruns_dir: str = "",
    experiment_name: str = "",
    rational: Optional[str] = None,
    details: Optional[str] = None,
    status: Optional[str] = None,
) -> Dict[str, Any]:
    row = _get_row(exp_id)
    if not row:
        raise ValueError(f"experiment {exp_id} not found")

    fields: List[str] = []
    params: List[Any] = []
    status = status or "done"
    if status is None and row.get("status") in ("starting", "running"):
        status = "done"

    rational = (rational or row.get("rational") or "").strip()
    details = (details if details is not None else row.get("details") or "").strip()

    fields.append("rational = %s"); params.append(rational)
    fields.append("rational_embedding = %s"); params.append(_vector_literal(embed(rational)))
    fields.append("details = %s"); params.append(details)
    fields.append("details_embedding = %s"); params.append(_vector_literal(embed(details)) if details else None)
    if evaluation is not None:
        fields.append("evaluation = %s"); params.append(evaluation.strip())
    if metrics is not None:
        fields.append("metrics = %s"); params.append(json.dumps(json.loads(metrics)))
    if ref_id:
        fields.append("experiment_ref_id = %s"); params.append(ref_id.strip())
    if mlruns_dir:
        fields.append("mlruns_dir = %s"); params.append(mlruns_dir.strip())
    if experiment_name:
        fields.append("experiment_name = %s"); params.append(experiment_name.strip())
    fields.append("status = %s"); params.append(status)
    fields.append("end_ts = %s"); params.append(_now_iso())
    fields.append("updated_at = %s"); params.append(_now_iso())
    params.append(exp_id)

    with _conn() as conn, conn.cursor() as cur:
        cur.execute(f"UPDATE rd_experiments SET {', '.join(fields)} WHERE id = %s", params)
        conn.commit()
    return _row_json(_get_row(exp_id) or {})


def _mlruns_dir(exp_name: str) -> str:
    uri = (os.environ.get("DATABASE_URL") or "").strip()
    if uri.startswith("postgres://"):
        uri = "postgresql+psycopg://" + uri[len("postgres://") :]
    if uri.startswith("postgresql://") or uri.startswith("postgresql+psycopg://"):
        with _conn() as conn, conn.cursor() as cur:
            cur.execute("SELECT artifact_location FROM experiments WHERE name = %s", (exp_name,))
            row = cur.fetchone()
        if not row:
            raise RuntimeError(f"mlflow experiment {exp_name!r} not found")
        return row["artifact_location"]

    lake = (os.environ.get("TAC_LAKE_DIR") or "").strip()
    if not lake:
        raise RuntimeError("TAC_LAKE_DIR not set")
    db_path = Path(lake) / "mlruns.db"
    if not db_path.exists():
        raise RuntimeError(f"mlruns.db not found at {db_path}")
    conn = sqlite3.connect(db_path)
    try:
        row = conn.execute("SELECT artifact_location FROM experiments WHERE name = ?", (exp_name,)).fetchone()
    finally:
        conn.close()
    if not row:
        raise RuntimeError(f"mlflow experiment {exp_name!r} not found in {db_path}")
    return row[0]


# --------------------------------------------------------------------------- git
def _parent_root() -> Path:
    start = Path.cwd()
    dir = start
    while dir != Path(dir.anchor):
        if (dir / "pnpm-workspace.yaml").exists() or (dir / "Cargo.toml").exists() or (dir / "opencode.json").exists():
            return dir
        dir = dir.parent
    raise RuntimeError("not inside a tradeac workspace")


def _gitc(*args: str) -> subprocess.CompletedProcess:
    root = _parent_root()
    exp = root / "experiments"
    user = (os.environ.get("GIT_USER") or "").strip()
    password = (os.environ.get("GIT_PASS") or "").strip()
    helper = f'!f() {{ echo "username={user}"; echo "password={password}"; }}; f'
    cmd = ["git", "-C", str(exp), "-c", f"credential.helper={helper}"] + list(args)
    return subprocess.run(cmd, capture_output=True, text=True)


def _git_ok(proc: subprocess.CompletedProcess) -> bool:
    return proc.returncode == 0


def _git_out(proc: subprocess.CompletedProcess) -> str:
    return (proc.stdout or "").strip() or (proc.stderr or "").strip()


def _require_auth() -> None:
    if not (os.environ.get("GIT_REPO_URL") or "").strip() or not (os.environ.get("GIT_USER") or "").strip():
        raise RuntimeError("GIT_REPO_URL / GIT_USER not set")


def _ensure_repo() -> None:
    root = _parent_root()
    exp = root / "experiments"
    if (exp / ".git").exists():
        _gitc("remote", "set-url", "origin", os.environ["GIT_REPO_URL"])
    else:
        _require_auth()
        (root / "experiments").mkdir(parents=True, exist_ok=True)
        subprocess.run(
            ["git", "clone", "-q", os.environ["GIT_REPO_URL"], str(exp)],
            check=True, capture_output=True, text=True,
        )
    _gitc("config", "user.email", f"{os.environ.get('GIT_USER', '')}@tradeac.local")
    _gitc("config", "user.name", os.environ.get("GIT_USER", "tradeac-agent"))


def _ensure_base(base: str = "main") -> str:
    _require_auth()
    _gitc("fetch", "origin", base)
    if _git_ok(_gitc("rev-parse", "--verify", f"origin/{base}")):
        return f"origin/{base}"
    if _git_ok(_gitc("rev-parse", "--verify", base)):
        return base
    return base


def _fork_branch(base_ref: str, branch: str) -> str:
    if _git_ok(_gitc("rev-parse", "--verify", f"origin/{branch}")):
        _gitc("checkout", "-q", "-B", branch, f"origin/{branch}")
        _gitc("reset", "-q", "--hard", f"origin/{branch}")
        return "reused existing branch"
    _gitc("fetch", "-q", "origin")
    if _git_ok(_gitc("rev-parse", "--verify", f"origin/{branch}")):
        _gitc("checkout", "-q", "-B", branch, f"origin/{branch}")
        return "reused existing branch"
    base_commit = ""
    if _git_ok(_gitc("rev-parse", "--verify", f"{base_ref}^{{commit}}")):
        base_commit = base_ref
    elif not base_ref.startswith("origin/"):
        base_commit = f"origin/{base_ref}"
    if not base_commit:
        raise RuntimeError(f"base '{base_ref}' not found locally or on origin")
    _gitc("checkout", "-q", "-B", branch, base_commit)
    return f"forked from {base_ref}"


def _snapshot_code(paths: Optional[List[str]] = None) -> str:
    root = _parent_root()
    exp = root / "experiments"
    paths = paths or ["tac-qlib/tac_qlib/contrib", "tac-qlib/tac_qlib/data"]
    parent_head = _git_out(_gitc("rev-parse", "HEAD")) or "unknown"
    import shutil

    shutil.rmtree(exp / "code", ignore_errors=True)
    (exp / "code").mkdir(parents=True, exist_ok=True)
    manifest = [f"# TradeAC custom-qlib-code snapshot (auto-generated)", f"# parent repo HEAD : {parent_head}"]
    for p in paths:
        manifest.append(f"#   {p}")
    manifest.append("# per-file hashes (git hash-object):")
    for p in paths:
        src = root / p
        if not src.exists():
            continue
        dst = exp / "code" / p
        dst.parent.mkdir(parents=True, exist_ok=True)
        if src.is_dir():
            shutil.copytree(src, dst, dirs_exist_ok=True)
            for f in sorted(src.rglob("*")):
                if f.is_file():
                    rel = str(f.relative_to(root))
                    h = _git_out(_gitc("hash-object", str(f)))
                    manifest.append(f"  {h}  {rel}")
        else:
            shutil.copy2(src, dst)
            h = _git_out(_gitc("hash-object", str(src)))
            manifest.append(f"  {h}  {p}")
    (exp / "code" / "MANIFEST.txt").write_text("\n".join(manifest) + "\n")
    return f"code snapshotted -> experiments/code (parent @ {parent_head[:12]})"


def _commit(message: str) -> str:
    _gitc("add", "-A")
    if _git_ok(_gitc("diff", "--cached", "--quiet")):
        return "nothing to commit"
    _gitc("commit", "-q", "-m", message)
    return "committed"


def _commit_push(message: str) -> str:
    result = _commit(message)
    if result == "nothing to commit":
        return result
    branch = _git_out(_gitc("branch", "--show-current"))
    _require_auth()
    proc = _gitc("push", "-u", "origin", branch)
    if not _git_ok(proc):
        raise RuntimeError(f"push failed: {_git_out(proc)}")
    return f"pushed {branch}"


def _parent_changes() -> str:
    root = _parent_root()
    proc = subprocess.run(["git", "-C", str(root), "status", "--porcelain"], capture_output=True, text=True)
    out = (proc.stdout or "").strip()
    if not out:
        return "parent repo clean (no changes)"
    lines = out.splitlines()
    filtered = [
        l for l in lines
        if not l.startswith(".. experiments/")
        and not l.startswith("?? experiments/")
        and not l.startswith(".. tac-qlib/tac_qlib/contrib/")
        and not l.startswith(".. tac-qlib/tac_qlib/data/")
        and not l.startswith("?? tac-qlib/tac_qlib/contrib/")
        and not l.startswith("?? tac-qlib/tac_qlib/data/")
    ]
    expected = [l for l in lines if l.startswith(".. tac-qlib/tac_qlib/contrib/") or l.startswith(".. tac-qlib/tac_qlib/data/")]
    note = ""
    if expected:
        note = "note: custom qlib code changed in the parent repo (contrib/data) — snapshotted to the experiment branch via trace snapshot:\n" + "\n".join(expected)
    if not filtered:
        return "parent repo changes limited to experiments/ and snapshotted custom qlib code (ok)" + (f"\n{note}" if note else "")
    return "WARNING: unexpected parent-repo changes outside the experiments/ clone:\n" + "\n".join(filtered) + "\n→ review and revert before finishing" + (f"\n{note}" if note else "")


def slugify(text: str) -> str:
    s = "".join(c for c in text.lower() if c.isalnum() or c in " -").replace(" ", "-")
    return s[:40].strip("-")


def get_experiment_branch(exp_id: int) -> str:
    row = _get_row(exp_id)
    if not row:
        raise ValueError(f"experiment {exp_id} not found")
    return row["git_branch"] or f"exp/{exp_id}"


# --------------------------------------------------------------------------- MCP tools
def rd_trace_init() -> dict:
    """Ensure the traceability store + experiments git repo are ready (rd_experiments table, base main)."""
    _init_db()
    _ensure_repo()
    base = _ensure_base("main")
    return {"status": "ready", "base": base}


def rd_trace_start(
    rational: str,
    details: str = "",
    experiment_name: str = "",
    evolved_from: str = "none",
    session_id: str = "",
) -> dict:
    """Open a traced experiment: insert the rd_experiments row, resolve evolved_from, fork + push the experiment branch. Pass `session_id` (the opencode chat id) so the lineage keeps a stable chat link. Returns experiment_id / branch / evolved_from / base_branch as one JSON object."""
    _init_db()
    _ensure_repo()

    evo_id = evolved_from
    if evolved_from == "auto":
        evo_id = str(_search_evolved_from(rational) or "")

    row = _start(rational, details, evolved_from=evo_id or "none", experiment_name=experiment_name, session_id=session_id)
    exp_id = int(row["id"])
    branch = f"exp/{exp_id}-{slugify(rational)}"
    _gitc("checkout", "-q", "-B", branch, "main") if False else None

    # fork from the predecessor branch (or main)
    base_branch = "main"
    if evo_id and evo_id.isdigit():
        base_branch = get_experiment_branch(int(evo_id))
    base_ref = _ensure_base(base_branch)
    _fork_branch(base_ref, branch)
    _snapshot_code()
    _gitc("checkout", "-q", "-B", branch, branch) if False else None

    # persist branch on the row
    with _conn() as conn, conn.cursor() as cur:
        cur.execute("UPDATE rd_experiments SET git_branch = %s WHERE id = %s", (branch, exp_id))
        conn.commit()

    _commit_push(f"start experiment {exp_id} ({branch})")
    return {"experiment_id": exp_id, "branch": branch, "evolved_from": evo_id or "none", "base_branch": base_branch}


def rd_trace_finish(
    experiment_id: int,
    ref_id: str = "",
    evaluation: str = "",
    metrics: str = "",
    mlruns_dir: str = "",
    experiment_name: str = "",
) -> dict:
    """Close a traced experiment: update the row (link the mlflow run, metrics/evaluation/end_ts), snapshot code, commit + push the branch. Returns the updated row."""
    _finish(experiment_id, ref_id=ref_id, evaluation=evaluation or None, metrics=metrics or None, mlruns_dir=mlruns_dir, experiment_name=experiment_name)
    branch = get_experiment_branch(experiment_id)
    _gitc("checkout", "-q", "-B", branch, branch)
    _snapshot_code()
    _commit_push(f"finish experiment {experiment_id} ({branch})")
    return {"experiment_id": experiment_id, "branch": branch, "row": _row_json(_get_row(experiment_id) or {})}


def rd_trace_commit(experiment_id: int, message: str = "wip") -> dict:
    """Commit the current experiment branch state (no push)."""
    branch = get_experiment_branch(experiment_id)
    _gitc("checkout", "-q", "-B", branch, branch)
    result = _commit(f"exp {experiment_id}: {message}")
    return {"experiment_id": experiment_id, "branch": branch, "result": result}


def rd_trace_snapshot(experiment_id: int, paths: str = "") -> dict:
    """Snapshot custom qlib contrib/data code onto the experiment branch (default contrib+data)."""
    branch = get_experiment_branch(experiment_id)
    _gitc("checkout", "-q", "-B", branch, branch)
    path_list = [p.strip() for p in paths.split(",") if p.strip()] if paths else None
    msg = _snapshot_code(path_list)
    _commit_push(f"exp {experiment_id}: snapshot custom qlib code")
    return {"experiment_id": experiment_id, "branch": branch, "result": msg}


def rd_trace_guard() -> dict:
    """Check the parent repo for unexpected changes outside the experiments clone."""
    return {"parent_changes": _parent_changes()}


def rd_trace_search(query: str, limit: int = 10, min_score: float = _MIN_SCORE) -> dict:
    """Semantic search over experiment rationals/details (pgvector, falls back to ILIKE)."""
    return {"results": _search(query, limit=limit, min_score=min_score)}


def rd_trace_get(experiment_id: int) -> dict:
    """Return one traced experiment row."""
    row = _get_row(experiment_id)
    if not row:
        raise ValueError(f"experiment {experiment_id} not found")
    return _row_json(row)


def rd_trace_list(limit: int = 20) -> dict:
    """List traced experiments (newest first)."""
    return {"experiments": [_row_json(r) for r in _all_rows(limit)]}


def rd_trace_mlruns_dir(experiment_name: str) -> dict:
    """Resolve the mlflow artifact location for an experiment name."""
    return {"mlruns_dir": _mlruns_dir(experiment_name)}


def register_trace_tools(server) -> None:
    """Attach all rd_trace_* tools to an MCPServer instance (called by rd_server.main())."""
    for fn in (
        rd_trace_init,
        rd_trace_start,
        rd_trace_finish,
        rd_trace_commit,
        rd_trace_snapshot,
        rd_trace_guard,
        rd_trace_search,
        rd_trace_get,
        rd_trace_list,
        rd_trace_mlruns_dir,
    ):
        server.tool(structured_output=False)(fn)