#!/usr/bin/env python3
"""Regenerate the model-registry artifacts from the single source of truth.

``.opencode/models/models.yaml`` is the ONE file humans edit to assign an
LLM model to an agent. This script regenerates every artifact derived from
it so a model assignment can never drift across surfaces:

  - ``.opencode/models/<agent>.txt`` — one per non-default agent; read by
    the dispatcher's ``_opencode_worker._resolve_role_model``.
  - ``.opencode/models/default.txt`` — the fallback model.
  - the ``agent`` block of ``opencode.json`` — each agent points at a
    ``{file:.../<agent>.txt}`` model file.
  - the ``model:`` frontmatter line of every managed agent's
    ``.opencode/agents/<agent>.md``.

opencode.json and the ``.md`` frontmatter are BOTH written, because the
repo's own history disagrees on which one OpenCode honours at startup —
writing both makes the assignment correct regardless. All four surfaces
are generated from the one manifest, so they cannot disagree.

A per-agent ``.txt`` whose agent is no longer in the manifest is an
ORPHAN: the dispatcher's resolver reads ``<agent>.txt`` before
``default.txt``, so a leftover file would silently pin the old model.
This script deletes such orphans on every run (tier ``.txt`` files
excepted — those belong to ``sync_tier_models.py``).

Tier-ladder agents (``task-implementor-tier-*``) are governed separately
by ``tiers.yaml`` + ``tools/sync_tier_models.py``; this script passes
their opencode.json entries through untouched and never edits their
``.md`` files. Agents with no opencode.json ``agent`` entry (e.g.
``task-implementor``) deliberately inherit their caller's model and are
left alone.

Usage:
    python3 tools/sync_models.py            # regenerate (idempotent)
    python3 tools/sync_models.py --check    # verify only; exit 1 on drift

The pipeline launcher runs this with no args at every startup, before
OpenCode caches its config, so the registry is always applied.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

import yaml

REPO_ROOT = Path(__file__).resolve().parent.parent
MODELS_DIR = REPO_ROOT / ".opencode" / "models"
AGENTS_DIR = REPO_ROOT / ".opencode" / "agents"
MODELS_YAML = MODELS_DIR / "models.yaml"
OPENCODE_JSON = REPO_ROOT / "opencode.json"

# Tier variants are owned by tiers.yaml + sync_tier_models.py — their
# opencode.json entries and .md files are left exactly as-is here.
TIER_AGENT_PREFIX = "task-implementor-tier-"

# Frontmatter of an agent .md: '---\n' ... '\n' '---\n'. Non-greedy so it
# stops at the FIRST closing delimiter (not a later markdown '---' rule).
_FRONTMATTER_RE = re.compile(r"^---\n(.*?\n)---\n", re.DOTALL)
# A top-level `model:` line WITHIN the frontmatter (no leading indent, so
# a nested key or a prose mention in the body is never matched).
_MODEL_LINE_RE = re.compile(r"^model:.*$", re.MULTILINE)


class SyncError(RuntimeError):
    """A manifest / validation problem that should abort the sync."""


def _load_manifest() -> tuple[str, dict[str, str]]:
    """Parse models.yaml → (default_model, {agent: model})."""
    if not MODELS_YAML.exists():
        raise SyncError(f"manifest not found: {MODELS_YAML}")
    data = yaml.safe_load(MODELS_YAML.read_text(encoding="utf-8")) or {}
    default = data.get("default")
    agents = data.get("agents") or {}
    if not isinstance(default, str) or "/" not in default:
        raise SyncError("models.yaml: `default` must be a 'providerID/modelID' string")
    if not isinstance(agents, dict):
        raise SyncError("models.yaml: `agents` must be a mapping")
    for name, model in agents.items():
        if not isinstance(model, str) or "/" not in model:
            raise SyncError(
                f"models.yaml: agents.{name} must be a 'providerID/modelID' string"
            )
        if name.startswith(TIER_AGENT_PREFIX):
            raise SyncError(
                f"models.yaml: agents.{name} is a tier agent — assign it in "
                "tiers.yaml, not here"
            )
    return default, dict(agents)


def _validate_providers(default: str, agents: dict[str, str], opencode: dict) -> None:
    """Every model's provider half must be declared in opencode.json."""
    providers = set(opencode.get("provider", {}))
    for label, model in [("default", default), *agents.items()]:
        provider = model.split("/", 1)[0]
        if provider not in providers:
            raise SyncError(
                f"model for {label!r} uses provider {provider!r}, which is not "
                "declared in opencode.json's `provider` block"
            )


def _expected_txt_files(default: str, agents: dict[str, str]) -> dict[Path, str]:
    """The .txt files this registry owns, mapped to their exact content."""
    out = {MODELS_DIR / "default.txt": default + "\n"}
    for name, model in agents.items():
        out[MODELS_DIR / f"{name}.txt"] = model + "\n"
    return out


def _agent_order(current_agents: dict, agents: dict[str, str]) -> list[str]:
    """The opencode.json ``agent`` block's key order after the manifest
    is applied: existing agents keep their file order; manifest agents
    not already present are appended in manifest order.

    Both the rendered ``agent`` block and the set of ``.md`` files to
    manage derive from this one ordering, so the two channels can never
    end up covering different agents.
    """
    keys = list(current_agents.keys())
    for name in agents:
        if name not in keys:
            keys.append(name)
    return keys


def _render_agent_block(
    keys: list[str], current_agents: dict, agents: dict[str, str]
) -> str:
    """Render the opencode.json ``agent`` block for the ordered ``keys``.

    Non-tier agents point at a ``{file:...}`` model file (their own when
    overridden, else ``default.txt``); tier agents are passed through
    with their current opencode.json value.
    """
    resolved: dict[str, str] = {}
    for name in keys:
        if name.startswith(TIER_AGENT_PREFIX):
            resolved[name] = str(current_agents[name].get("model", ""))
        elif name in agents:
            resolved[name] = f"{{file:./.opencode/models/{name}.txt}}"
        else:
            resolved[name] = "{file:./.opencode/models/default.txt}"

    width = max(len(f'"{k}":') for k in keys)
    lines = ['    "agent": {']
    for i, name in enumerate(keys):
        comma = "" if i == len(keys) - 1 else ","
        label = f'"{name}":'.ljust(width)
        lines.append(
            "        " + label + ' { "model": "' + resolved[name] + '" }' + comma
        )
    lines.append("    }")
    return "\n".join(lines)


def _splice_opencode_json(block: str) -> str:
    """Return opencode.json text with its ``agent`` block replaced.

    The ``agent`` key is the last top-level key, so everything from
    ``\\n    "agent": {`` to EOF is regenerated; the rest of the file is
    preserved byte-for-byte.
    """
    text = OPENCODE_JSON.read_text(encoding="utf-8")
    match = re.search(r'\n    "agent": \{', text)
    if not match:
        raise SyncError('could not locate the "agent" block in opencode.json')
    return text[: match.start()] + "\n" + block + "\n}\n"


def _md_with_model(text: str, model: str, label: str) -> str:
    """Return the agent-.md ``text`` with its frontmatter ``model:`` line
    set to ``model`` — replacing an existing line in place, or inserting
    one as the first frontmatter key. Everything else is byte-preserved.
    """
    match = _FRONTMATTER_RE.match(text)
    if not match:
        raise SyncError(f"{label}: no YAML frontmatter ('--- ... ---')")
    frontmatter = match.group(1)
    desired = f"model: {model}"
    if _MODEL_LINE_RE.search(frontmatter):
        new_fm = _MODEL_LINE_RE.sub(lambda _m: desired, frontmatter, count=1)
    else:
        new_fm = desired + "\n" + frontmatter
    return "---\n" + new_fm + "---\n" + text[match.end() :]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Sync the model registry from .opencode/models/models.yaml"
    )
    parser.add_argument(
        "--check",
        action="store_true",
        help="verify the generated artifacts are in sync; write nothing; "
        "exit 1 on drift",
    )
    args = parser.parse_args()

    try:
        default, agents = _load_manifest()
        opencode = json.loads(OPENCODE_JSON.read_text(encoding="utf-8"))
        _validate_providers(default, agents, opencode)
        if not isinstance(opencode.get("agent"), dict):
            raise SyncError("opencode.json has no `agent` block")
        current_agents = opencode["agent"]

        # Every manifest agent must name a REAL agent — i.e. have an
        # `.opencode/agents/<name>.md`. Without this, a typo'd manifest
        # key is silently appended to opencode.json as a bogus agent
        # (and gets an orphan .txt), while the real agent it was meant
        # to name keeps the default model — no error, wrong result.
        for name in agents:
            if not (AGENTS_DIR / f"{name}.md").exists():
                raise SyncError(
                    f"models.yaml: agents.{name} has no "
                    f".opencode/agents/{name}.md — a manifest key must name a "
                    "real agent (check for a typo)"
                )

        # One ordering drives BOTH OpenCode-facing channels (the
        # opencode.json block and the .md frontmatter), so a manifest
        # agent appended to opencode.json also gets its .md written.
        agent_order = _agent_order(current_agents, agents)

        txt_files = _expected_txt_files(default, agents)
        expected_json = _splice_opencode_json(
            _render_agent_block(agent_order, current_agents, agents)
        )
        # Validate the regenerated JSON BEFORE writing anything, so a
        # bad render aborts cleanly instead of leaving opencode.json
        # corrupt on disk.
        try:
            json.loads(expected_json)
        except ValueError as exc:
            raise SyncError(
                f"regenerated opencode.json is not valid JSON ({exc})"
            ) from exc

        # Every non-tier agent in the final order gets its .md
        # frontmatter `model:` set too (the "both channels" guarantee).
        md_files: dict[Path, str] = {}
        for name in agent_order:
            if name.startswith(TIER_AGENT_PREFIX):
                continue
            md_path = AGENTS_DIR / f"{name}.md"
            if not md_path.exists():
                raise SyncError(
                    f"agent {name!r} is in opencode.json but has no {md_path}"
                )
            model = agents.get(name, default)
            md_files[md_path] = _md_with_model(
                md_path.read_text(encoding="utf-8"), model, md_path.name
            )

        # Orphan .txt files: a non-tier <agent>.txt left behind after
        # its override was dropped from models.yaml would still be read
        # FIRST by the dispatcher's resolver (<agent>.txt before
        # default.txt), silently pinning the stale model. Tier .txt
        # files belong to sync_tier_models.py and are never touched.
        expected_txt = set(txt_files)
        orphan_txt = sorted(
            p
            for p in MODELS_DIR.glob("*.txt")
            if p not in expected_txt and not p.stem.startswith(TIER_AGENT_PREFIX)
        )
    except (SyncError, ValueError) as exc:
        print(f"sync_models: ERROR: {exc}", file=sys.stderr)
        return 1

    rel = lambda p: str(p.relative_to(REPO_ROOT))  # noqa: E731
    drift: list[Path] = []
    for path, content in txt_files.items():
        if not path.exists() or path.read_text(encoding="utf-8") != content:
            drift.append(path)
    if OPENCODE_JSON.read_text(encoding="utf-8") != expected_json:
        drift.append(OPENCODE_JSON)
    for path, content in md_files.items():
        if path.read_text(encoding="utf-8") != content:
            drift.append(path)
    drift.extend(orphan_txt)

    if args.check:
        if drift:
            print(
                f"sync_models: DRIFT — {len(drift)} artifact(s) out of sync with "
                "models.yaml:\n  " + "\n  ".join(sorted(rel(p) for p in drift)),
                file=sys.stderr,
            )
            return 1
        print("sync_models: registry in sync with models.yaml")
        return 0

    if not drift:
        print("sync_models: registry already in sync — no changes")
        return 0

    for path, content in txt_files.items():
        path.write_text(content, encoding="utf-8")
    OPENCODE_JSON.write_text(expected_json, encoding="utf-8")
    for path, content in md_files.items():
        path.write_text(content, encoding="utf-8")
    for path in orphan_txt:
        path.unlink()

    orphan_set = set(orphan_txt)
    updated = sorted(rel(p) for p in drift if p not in orphan_set)
    parts: list[str] = []
    if updated:
        md_changed = sum(1 for p in updated if p.endswith(".md"))
        listed = [p for p in updated if not p.endswith(".md")]
        if md_changed:
            listed.append(f"{md_changed} agent .md file(s)")
        parts.append("updated " + ", ".join(listed))
    if orphan_txt:
        parts.append(
            "removed orphan "
            + ", ".join(sorted(rel(p) for p in orphan_txt))
            + " (no longer in models.yaml)"
        )
    print("sync_models: applied registry — " + "; ".join(parts))
    return 0


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