"""
install_hook.py  —  stlc_agents.shared.install_hook
─────────────────────────────────────────────────────
Called automatically by postinstall.js after `pip install qa-gentic-stlc-agents`.
Also callable manually: python -m stlc_agents.shared.install_hook

Applies the cost tracking patch to all MCP server files by importing
and running the same logic as scripts/apply_cost_tracking.py, but resolved
relative to the installed package location (works in site-packages, .venv, etc).
"""

from __future__ import annotations
import re
import sys
from pathlib import Path


SERVERS = [
    ("agent_gherkin_generator",     "qa-gherkin-generator"),
    ("agent_test_case_manager",     "qa-test-case-manager"),
    ("agent_playwright_generator",  "qa-playwright-generator"),
    ("agent_helix_writer",          "qa-helix-writer"),
    ("agent_jira_manager",          "qa-jira-manager"),
    ("agent_migration",             "qa-migration"),
]

IMPORT_MARKER = "from stlc_agents.shared.cost_tracker import track"
TIME_IMPORT   = "import time"

OLD_RETURN = (
    'return [types.TextContent(type="text", '
    'text=json.dumps(result, indent=2, ensure_ascii=False))]'
)
NEW_RETURN = "return track(result, tool_name=name, server={server!r}, t0=t0)"

OLD_ERR_BLOCK = """\
        return [types.TextContent(
            type="text",
            text=json.dumps({"error": str(exc), "tool": name}, indent=2),
        )]"""

NEW_ERR_BLOCK = """\
        err_result = {{"error": str(exc), "tool": name}}
        return track(err_result, tool_name=name, server={server!r}, t0=t0)"""


def _root() -> Path:
    """Resolve the stlc_agents package root regardless of install method."""
    import stlc_agents
    return Path(stlc_agents.__file__).parent


def patch_server(agent_dir: str, server_name: str, root: Path) -> str:
    """Patch one server file. Returns 'patched' | 'already_patched' | 'not_found' | 'no_change'."""
    path = root / agent_dir / "server.py"
    if not path.exists():
        return "not_found"

    src = path.read_text(encoding="utf-8")
    if IMPORT_MARKER in src:
        return "already_patched"

    original = src

    # 1. import time
    if TIME_IMPORT not in src:
        src = src.replace("import sys\n", "import sys\nimport time\n", 1)

    # 2. cost_tracker import — after last `from stlc_agents...` line
    last_match = None
    for m in re.finditer(r"^from stlc_agents\..+\n", src, re.MULTILINE):
        last_match = m
    if last_match:
        pos = last_match.end()
        src = src[:pos] + "from stlc_agents.shared.cost_tracker import track\n" + src[pos:]
    else:
        src = src.replace(
            "from dotenv import load_dotenv\n",
            "from dotenv import load_dotenv\nfrom stlc_agents.shared.cost_tracker import track\n",
            1,
        )

    # 3. t0 = time.monotonic() inside call_tool()
    src = src.replace(
        "@app.call_tool()\nasync def call_tool(name: str, arguments: dict)"
        " -> list[types.TextContent]:\n    try:\n",
        "@app.call_tool()\nasync def call_tool(name: str, arguments: dict)"
        " -> list[types.TextContent]:\n    t0 = time.monotonic()\n    try:\n",
        1,
    )

    # 4. Replace all result return lines
    new_ret = NEW_RETURN.format(server=server_name)
    for indent in ("        ", "            ", "                "):
        src = src.replace(f"{indent}{OLD_RETURN}", f"{indent}{new_ret}")

    # 5. Replace error-path block
    src = src.replace(OLD_ERR_BLOCK, NEW_ERR_BLOCK.format(server=server_name))

    if src == original:
        return "no_change"

    path.write_text(src, encoding="utf-8")
    return "patched"


def apply_cost_tracking() -> None:
    """Entry point — called by postinstall.js and the console_script."""
    root = _root()

    ok   = "\x1b[32m✓\x1b[0m"
    skip = "\x1b[33m–\x1b[0m"
    err  = "\x1b[31m✗\x1b[0m"

    print("\n  stlc-agents · Activating cost tracking on MCP servers...\n")

    any_patched = False
    for agent_dir, server_name in SERVERS:
        status = patch_server(agent_dir, server_name, root)
        if status == "patched":
            print(f"  {ok}  {agent_dir}  ({server_name})")
            any_patched = True
        elif status == "already_patched":
            print(f"  {skip}  {agent_dir}  — already active")
        elif status == "not_found":
            print(f"  {err}  {agent_dir}/server.py  — not found (skip)")
        elif status == "no_change":
            print(f"  {skip}  {agent_dir}  — no matching pattern (manual patch needed)")

    print()
    if any_patched:
        print("  Cost tracking is now active.  On every MCP tool call you will see:")
        print("  [stlc-cost] <server> · <tool>  ~<N>K tokens  $<cost>  (session: $<total>)")
        print()
        print("  Session logs:  ~/.qa-stlc/cost-<session-id>.jsonl")
        print("  View report:   qa-stlc cost")
        print("  View all:      qa-stlc cost --all")
        print()
        print("  Environment variables:")
        print("    STLC_COST_TRACKING=false          disable output")
        print("    STLC_CODING_AGENT_MODEL=<model>   set your agent's model for exact pricing")
        print("      e.g. claude-sonnet-4-6 | claude-opus-4-6 | gpt-4o")
        print("    STLC_COST_LOG_DIR=<path>          change log directory")
    else:
        print("  All servers already have cost tracking active.")

    print()


def main() -> None:
    apply_cost_tracking()


if __name__ == "__main__":
    main()