"""
cost_tracker.py  —  stlc_agents.shared.cost_tracker
─────────────────────────────────────────────────────
Three tracking modes:

  track()                       — MCP server tool calls (coding-agent-driven flow).
                                  Token counts are ESTIMATED from payload size because
                                  the MCP server subprocess never sees the coding agent's
                                  API response.

  track_llm_call()              — Webhook orchestrator LLM calls (agent_runner.py).
                                  Token counts are EXACT: taken directly from the LLM
                                  API response's usage block, matching how promptfoo
                                  tracks costs (input + output + cache per iteration,
                                  accumulated across all iterations in the agent loop).

  _track_parent_agent_atexit()  — Claude Code parent-agent reasoning cost.
                                  Reads the session transcript JSONL Claude Code
                                  writes at `~/.claude/projects/<encoded-cwd>/<session>.jsonl`
                                  and sums the EXACT `message.usage` block on every
                                  assistant turn. This captures both the agent's own
                                  reasoning AND the tokens it spent on responses from
                                  every other MCP (Playwright, Atlassian, etc.) —
                                  those bytes are already in the parent's `input_tokens`
                                  on the next turn, so a single transcript parse covers
                                  everything the agent actually paid for. Run from atexit
                                  so the transcript is fully flushed by the time we read it.

MODEL AUTO-DETECTION (for track() only)
─────────────────────────────────────────
  Agent              Env var set automatically         Value example
  ─────────────────  ──────────────────────────────   ──────────────────────────
  Claude Code        ANTHROPIC_MODEL                  claude-sonnet-4-6
                     CLAUDE_MODEL (fallback)           claude-opus-4-6
  GitHub Copilot     GITHUB_COPILOT_MODEL (if set)    gpt-4o
  Cursor / Windsurf  (none — user sets manually)       —
  Any agent          STLC_CODING_AGENT_MODEL           user-specified override

Detection order (first match wins):
  1. STLC_CODING_AGENT_MODEL   — explicit user override (always wins)
  2. ANTHROPIC_MODEL            — set by Claude Code automatically
  3. CLAUDE_MODEL               — older Claude Code versions
  4. GITHUB_COPILOT_MODEL       — Copilot if configured
  5. ~/.qa-stlc/agent-model     — saved preference from `qa-stlc cost --set-model`
  6. "claude-sonnet-4-6"        — safe default (most common)

WHAT GETS LOGGED (per tool call / LLM call)
─────────────────────────────────────────────
  ~/.qa-stlc/cost-<session>.jsonl — machine-readable session log.
  stderr live line — visible in Claude Code MCP log, VS Code Output, etc.
  atexit summary — printed when the process exits.
"""

from __future__ import annotations

import atexit
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from mcp import types

from .pricing import ModelPricing, get_pricing

# ── Model detection ────────────────────────────────────────────────────────

_PREF_FILE = Path.home() / ".qa-stlc" / "agent-model"


def _detect_model() -> str:
    """
    Detect the coding agent's model from environment variables,
    in priority order. Falls back to saved preference, then default.
    """
    # 1. Explicit user override — always wins
    if v := os.getenv("STLC_CODING_AGENT_MODEL", "").strip():
        return v

    # 2. Claude Code sets this automatically when it spawns MCP subprocesses
    if v := os.getenv("ANTHROPIC_MODEL", "").strip():
        return v

    # 3. Older Claude Code versions used this name
    if v := os.getenv("CLAUDE_MODEL", "").strip():
        return v

    # 4. GitHub Copilot (if the MCP config passes it through)
    if v := os.getenv("GITHUB_COPILOT_MODEL", "").strip():
        return v

    # 5. Saved preference from `qa-stlc cost --set-model`
    try:
        if _PREF_FILE.exists():
            saved = _PREF_FILE.read_text(encoding="utf-8").strip()
            if saved:
                return saved
    except OSError:
        pass

    # 6. Safe default
    return "claude-sonnet-4-6"


# ── Config ─────────────────────────────────────────────────────────────────

_TRACKING_ENABLED = os.getenv("STLC_COST_TRACKING", "true").lower() != "false"
_LOG_DIR          = Path(os.getenv("STLC_COST_LOG_DIR", str(Path.home() / ".qa-stlc")))
_SESSION_ID       = os.getenv("STLC_SESSION_ID", f"stlc-{int(time.time())}")

# Resolved once at import time
_MODEL_ID         = _detect_model()
_PRICING          = get_pricing(_MODEL_ID)

# ANSI colours (suppressed when not a TTY, e.g. in VS Code Output pane)
_TTY = sys.stderr.isatty()
_C = {k: v if _TTY else "" for k, v in {
    "reset": "\x1b[0m", "bold": "\x1b[1m", "dim": "\x1b[2m",
    "cyan": "\x1b[36m", "green": "\x1b[32m", "yellow": "\x1b[33m",
}.items()}


# ── Session ────────────────────────────────────────────────────────────────

class _Session:
    def __init__(self):
        self.id         = _SESSION_ID
        self.started_at = time.time()
        self.records: list[dict] = []
        _LOG_DIR.mkdir(parents=True, exist_ok=True)
        self.log_path = _LOG_DIR / f"cost-{self.id}.jsonl"

    def add(self, record: dict) -> None:
        self.records.append(record)
        try:
            with self.log_path.open("a", encoding="utf-8") as f:
                f.write(json.dumps(record) + "\n")
        except OSError:
            pass  # never break a tool call over logging

    def running_total(self) -> float:
        return sum(r.get("cost_usd", 0.0) for r in self.records)


_session: Optional[_Session] = None


def _get_session() -> _Session:
    global _session
    if _session is None:
        _session = _Session()
    return _session


# ── Token estimation ───────────────────────────────────────────────────────

def _estimate_tokens(payload: Any) -> tuple[int, int, int]:
    """
    Estimate input/output tokens from the ADO/Jira response payload size.

    The MCP server cannot see the LLM's token usage. We estimate from the
    JSON payload the server returns, since the coding agent must tokenise
    that entire payload to read it:

      total  ≈ len(json_text) / 4     (chars-per-token heuristic)
      input  ≈ total * 0.70           (agent reading the ADO/Jira result)
      output ≈ total * 0.30           (agent writing the next artifact)

    Returns (total, input, output).
    """
    text  = json.dumps(payload) if not isinstance(payload, str) else payload
    total = max(1, len(text) // 4)
    return total, int(total * 0.70), int(total * 0.30)


# ── Public API ─────────────────────────────────────────────────────────────

def track(
    result: Any,
    *,
    tool_name: str,
    server: str,
    t0: float,
) -> list[types.TextContent]:
    """
    Wrap a call_tool() result. Injects _cost into the response JSON,
    logs the call, and returns the TextContent list ready for the MCP client.

    Called from each server's call_tool() as the final return:

        return track(result, tool_name=name, server="qa-gherkin-generator", t0=t0)

    Args:
        result:    The dict your tool function produced.
        tool_name: The MCP tool name, e.g. "fetch_work_item_for_gherkin".
        server:    The MCP server name, e.g. "qa-gherkin-generator".
        t0:        time.monotonic() captured at the start of call_tool().
    """
    if not _TRACKING_ENABLED:
        text = json.dumps(result, indent=2, ensure_ascii=False)
        return [types.TextContent(type="text", text=text)]

    latency_ms         = int((time.monotonic() - t0) * 1000)
    sess               = _get_session()
    total, inp, out    = _estimate_tokens(result)

    cost = _PRICING.cost(input_tokens=inp, output_tokens=out) if _PRICING else 0.0
    running = sess.running_total() + cost

    cost_block = {
        "session_id":       sess.id,
        "server":           server,
        "tool":             tool_name,
        "model":            _MODEL_ID,
        "model_source":     _model_source(),
        "input_tokens":     inp,
        "output_tokens":    out,
        "estimated_tokens": total,
        "cost_usd":         round(cost, 8),
        "latency_ms":       latency_ms,
        "timestamp":        datetime.now(timezone.utc).isoformat(),
        "session_total_usd": round(running, 8),
        "token_method":     "estimated",
        "token_note": (
            "Estimated from ADO/Jira payload size (chars÷4, 70/30 split). "
            "The MCP server has no access to the coding agent's token usage. "
            "Set STLC_CODING_AGENT_MODEL if the detected model is wrong. "
            f"Detected via: {_model_source()}."
        ),
    }

    if isinstance(result, dict):
        result["_cost"] = cost_block

    sess.add({"tool": tool_name, "server": server, **cost_block})
    _print_live(server, tool_name, total, cost, latency_ms, running)

    return [types.TextContent(
        type="text",
        text=json.dumps(result, indent=2, ensure_ascii=False),
    )]


def track_llm_call(
    *,
    model: str,
    provider: str,
    input_tokens: int,
    output_tokens: int,
    cache_write_tokens: int = 0,
    cache_read_tokens: int = 0,
    tool: str = "llm-agent-loop",
    server: str = "agent-runner",
    work_item_id: str = "",
    iterations: int = 1,
    latency_ms: int = 0,
) -> float:
    """
    Record exact LLM token usage from an API response (webhook / agent_runner path).

    Token counts come directly from the LLM API response.usage block —
    same approach as promptfoo: capture per-iteration, accumulate across all
    iterations, compute cost once at the end.

    Anthropic fields:
      input_tokens, output_tokens,
      cache_creation_input_tokens → cache_write_tokens,
      cache_read_input_tokens     → cache_read_tokens

    OpenAI fields:
      prompt_tokens  → input_tokens,
      completion_tokens → output_tokens,
      prompt_tokens_details.cached_tokens → cache_read_tokens

    Returns the cost in USD.
    """
    if not _TRACKING_ENABLED:
        return 0.0

    pricing = get_pricing(model)
    cost = (
        pricing.cost(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cache_write_tokens=cache_write_tokens,
            cache_read_tokens=cache_read_tokens,
        )
        if pricing else 0.0
    )

    sess    = _get_session()
    running = sess.running_total() + cost
    total   = input_tokens + output_tokens

    cache_note = ""
    if cache_write_tokens or cache_read_tokens:
        cache_note = f" cache_write={cache_write_tokens} cache_read={cache_read_tokens}"

    record = {
        "tool":               tool,
        "server":             server,
        "session_id":         sess.id,
        "model":              model,
        "model_source":       f"{provider}-api-response",
        "input_tokens":       input_tokens,
        "output_tokens":      output_tokens,
        "cache_write_tokens": cache_write_tokens,
        "cache_read_tokens":  cache_read_tokens,
        "estimated_tokens":   total,
        "cost_usd":           round(cost, 8),
        "latency_ms":         latency_ms,
        "timestamp":          datetime.now(timezone.utc).isoformat(),
        "session_total_usd":  round(running, 8),
        "token_method":       "exact",
        "token_note": (
            f"Exact counts from {provider} API response. "
            f"Iterations: {iterations}.{cache_note}"
        ),
        "work_item_id":       work_item_id,
        "iterations":         iterations,
    }
    sess.add(record)
    _print_live(server, tool, total, cost, latency_ms, running, exact=True)
    return cost


def get_session_id() -> str:
    """Return the current session ID (used to pass STLC_SESSION_ID to subprocesses)."""
    return _SESSION_ID


def track_healing(payload: dict) -> None:
    """
    Record an AI Vision healing call from LocatorHealer.ts.
    Called by healing_cost_server.py when LocatorHealer posts usage.
    """
    if not _TRACKING_ENABLED:
        return
    sess    = _get_session()
    model   = payload.get("model_id", "claude-sonnet-4-20250514")
    pricing = get_pricing(model)
    usage   = payload.get("usage", {})
    inp     = usage.get("input_tokens", usage.get("prompt_tokens", 0))
    out     = usage.get("output_tokens", usage.get("completion_tokens", 0))
    cost    = pricing.cost(input_tokens=inp, output_tokens=out) if pricing else 0.0

    record = {
        "tool":             f"ai-vision-{payload.get('provider', 'anthropic')}",
        "server":           "locator-healer",
        "session_id":       sess.id,
        "model":            model,
        "model_source":     "locator-healer-env",
        "input_tokens":     inp,
        "output_tokens":    out,
        "estimated_tokens": inp + out,
        "cost_usd":         round(cost, 8),
        "latency_ms":       payload.get("latency_ms", 0),
        "timestamp":        datetime.now(timezone.utc).isoformat(),
        "session_total_usd": round(sess.running_total() + cost, 8),
        "token_method":     "exact",
        "healing_meta": {
            "locator_key": payload.get("locator_key"),
            "healed":      payload.get("healed", False),
            "selector":    payload.get("selector"),
        },
    }
    sess.add(record)
    _print_live(
        "locator-healer",
        f"ai-vision ({payload.get('provider', '?')})",
        inp + out, cost, payload.get("latency_ms", 0),
        sess.running_total(),
    )


def _model_source() -> str:
    """Describe where the detected model came from."""
    if os.getenv("STLC_CODING_AGENT_MODEL", "").strip():
        return "STLC_CODING_AGENT_MODEL env var"
    if os.getenv("ANTHROPIC_MODEL", "").strip():
        return "ANTHROPIC_MODEL env var (set by Claude Code)"
    if os.getenv("CLAUDE_MODEL", "").strip():
        return "CLAUDE_MODEL env var (set by Claude Code)"
    if os.getenv("GITHUB_COPILOT_MODEL", "").strip():
        return "GITHUB_COPILOT_MODEL env var"
    if _PREF_FILE.exists():
        return f"saved preference ({_PREF_FILE})"
    return "default fallback"


# ── Live stderr line ───────────────────────────────────────────────────────

def _print_live(
    server: str, tool: str, tokens: int,
    cost: float, latency_ms: int, running: float,
    *, exact: bool = False,
) -> None:
    c        = _C
    prefix   = "" if exact else "~"
    tok_str  = f"{tokens/1000:.1f}K" if tokens >= 1000 else str(tokens)
    cost_str = f"${cost:.6f}"
    total    = f"${running:.6f}"
    method   = "" if exact else f"{c['dim']} [est]{c['reset']}"
    print(
        f"{c['dim']}[stlc-cost]{c['reset']} "
        f"{c['cyan']}{server}{c['reset']}{c['dim']} · {c['reset']}{tool}"
        f"  {prefix}{tok_str} tokens  {c['green']}{cost_str}{c['reset']}{method}"
        f"  {c['dim']}(session: {total}  {latency_ms}ms){c['reset']}",
        file=sys.stderr, flush=True,
    )


# ── Tool → artifact label ──────────────────────────────────────────────────

_TOOL_ARTIFACT: dict[str, str] = {
    "fetch_work_item_for_gherkin":  "Work item fetched",
    "fetch_feature_hierarchy":      "Feature hierarchy fetched",
    "generate_and_attach_gherkin":  "Gherkin attached to work item",
    "attach_gherkin_to_feature":    "Gherkin attached to feature",
    "attach_gherkin_to_work_item":  "Gherkin attached to work item",
    "validate_gherkin_content":     "Gherkin validated",
    "capture_app_context":          "App context captured",
    "generate_playwright_code":     "Playwright code generated",
    "scaffold_locator_repository":  "Locator repository scaffolded",
    "attach_code_to_work_item":     "Code attached to work item",
    "validate_gherkin_steps":       "Gherkin steps validated",
    "pre_validate_cucumber_steps":  "Cucumber steps pre-validated",
    "get_generated_files":          "Generated files retrieved",
    "inspect_helix_project":        "Helix project inspected",
    "write_helix_files":            "Files written to Helix",
    "update_helix_file":            "Helix file updated",
    "read_helix_file":              "Helix file read",
    "list_helix_tree":              "Helix tree listed",
    "fetch_work_item":              "Work item fetched",
    "create_and_link_test_cases":   "Test cases created & linked",
    "create_deduped_test_cases":    "Test cases created (deduped)",
    "get_linked_test_cases":        "Linked test cases retrieved",
    "llm-agent-loop":               "LLM orchestration",
}

_SERVER_ORDER = [
    "qa-test-case-manager", "qa-jira-manager",
    "qa-gherkin-generator", "qa-playwright-generator", "qa-helix-writer",
    "agent-runner",
]

_SERVER_FRIENDLY: dict[str, str] = {
    "qa-test-case-manager":   "QA Test Case Manager",
    "qa-jira-manager":        "QA Jira Manager",
    "qa-gherkin-generator":   "QA Gherkin Generator",
    "qa-playwright-generator": "QA Playwright Generator",
    "qa-helix-writer":        "QA Helix Writer",
    "agent-runner":           "Orchestrator",
    "locator-healer":         "Locator Healer",
}

_TOOL_FRIENDLY: dict[str, str] = {
    "fetch_work_item_for_gherkin":  "fetch WI",
    "fetch_feature_hierarchy":      "fetch feature hierarchy",
    "generate_and_attach_gherkin":  "generate & attach Gherkin",
    "attach_gherkin_to_feature":    "attach Gherkin to feature",
    "attach_gherkin_to_work_item":  "attach Gherkin to WI",
    "validate_gherkin_content":     "validate Gherkin",
    "capture_app_context":          "capture app context",
    "generate_playwright_code":     "generate Playwright code",
    "scaffold_locator_repository":  "scaffold locators",
    "attach_code_to_work_item":     "attach code to WI",
    "validate_gherkin_steps":       "validate Gherkin steps",
    "pre_validate_cucumber_steps":  "pre-validate Cucumber steps",
    "get_generated_files":          "retrieve generated files",
    "inspect_helix_project":        "inspect Helix project",
    "write_helix_files":            "write files to Helix",
    "update_helix_file":            "update Helix file",
    "read_helix_file":              "read Helix file",
    "list_helix_tree":              "list Helix tree",
    "fetch_work_item":              "fetch WI",
    "create_and_link_test_cases":   "create & link test cases",
    "create_deduped_test_cases":    "create test cases (deduped)",
    "get_linked_test_cases":        "get linked test cases",
    "llm-agent-loop":               "LLM agent loop",
    "ai-vision-anthropic":          "AI Vision (Anthropic)",
    "ai-vision-copilot":            "AI Vision (Copilot)",
}


def _step_label(record: dict) -> str:
    server = record.get("server", "?")
    tool   = record.get("tool", "?")
    sname  = _SERVER_FRIENDLY.get(server, server)
    tname  = _TOOL_FRIENDLY.get(tool, tool)
    if server == "agent-runner":
        return f"Orchestrator ({tname})"
    return f"{sname} ({tname})"


def _model_display(model_id: str) -> str:
    """Return a short display name for a model ID."""
    p = get_pricing(model_id)
    return p.display_name if p else model_id


def _fmt_tok(n: int, exact: bool) -> str:
    pfx = "" if exact else "~"
    if n >= 1_000_000:
        return f"{pfx}{n/1_000_000:.1f}M"
    if n >= 1000:
        return f"{pfx}{n/1000:.1f}K"
    return f"{pfx}{n}"


# ── Unified pipeline summary (called by agent_runner after loop) ───────────

def print_pipeline_summary(
    session_id: str,
    work_item_id: str = "",
    elapsed_s: float = 0.0,
    model: str = "",
    artefacts: list[dict] | None = None,
) -> None:
    """
    Print the two-section final report:
      1. Artefact Summary  — what was produced and where it lives
      2. Token and Cost Report — per-step token counts and USD cost
    """
    if not _TRACKING_ENABLED:
        return

    log_path = _LOG_DIR / f"cost-{session_id}.jsonl"
    records: list[dict] = []
    try:
        with log_path.open(encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError:
                        pass
    except OSError:
        return
    if not records:
        return

    c   = _C
    W   = 108
    m   = model or _MODEL_ID
    wi  = f"WI {work_item_id}  ·  " if work_item_id else ""
    dur = f"  ·  {elapsed_s:.1f}s" if elapsed_s else ""

    print(f"\n{c['bold']}{'═'*W}{c['reset']}", file=sys.stderr)
    print(
        f"{c['bold']}  stlc-agents · Pipeline Report  ·  {wi}{m}{dur}{c['reset']}",
        file=sys.stderr,
    )
    print(f"{c['bold']}{'═'*W}{c['reset']}", file=sys.stderr)

    # ── Section 1: Artefact Summary ──────────────────────────────────────────
    if artefacts:
        print(f"\n{c['bold']}  1. Artefact Summary{c['reset']}", file=sys.stderr)
        C1, C2, C3, C4 = 3, 26, 22, 22
        hdr = (
            f"  {'#':<{C1}}  {'Artefact':<{C2}}  {'Type':<{C3}}"
            f"  {'Location / Status':<{C4}}  Detail"
        )
        print(f"\n{hdr}", file=sys.stderr)
        print(f"  {'─'*(W-2)}", file=sys.stderr)
        for i, art in enumerate(artefacts, 1):
            row = (
                f"  {i:<{C1}}  {art.get('name',''):<{C2}}  "
                f"{art.get('type',''):<{C3}}  "
                f"{art.get('location',''):<{C4}}  "
                f"{c['dim']}{art.get('detail','')}{c['reset']}"
            )
            print(row, file=sys.stderr)

    # ── Section 2: Token and Cost Report ─────────────────────────────────────
    print(f"\n{c['bold']}  2. Token and Cost Report{c['reset']}", file=sys.stderr)
    S1, S2, S3, S4, S5 = 3, 46, 20, 12, 12
    hdr2 = (
        f"\n  {'#':<{S1}}  {'Agent / Step':<{S2}}  {'Model':<{S3}}"
        f"  {'Input':>{S4}}  {'Output':>{S5}}  {'Cost USD':>12}"
    )
    print(hdr2, file=sys.stderr)
    print(f"  {'─'*(W-2)}", file=sys.stderr)

    total_in = total_out = 0
    total_cost = 0.0

    for i, r in enumerate(records, 1):
        exact  = r.get("token_method") == "exact"
        inp    = r.get("input_tokens", 0)
        out    = r.get("output_tokens", 0)
        cost   = r.get("cost_usd", 0.0)
        mdl    = _model_display(r.get("model", m))
        label  = _step_label(r)
        pfx    = "" if exact else "~"
        in_s   = _fmt_tok(inp, exact)
        out_s  = _fmt_tok(out, exact)
        cost_s = f"{pfx}${cost:.6f}"

        print(
            f"  {i:<{S1}}  {label:<{S2}}  {mdl:<{S3}}"
            f"  {in_s:>{S4}}  {out_s:>{S5}}  "
            f"{c['green']}{cost_s:>12}{c['reset']}",
            file=sys.stderr,
        )

    # Session total: use only exact agent-runner rows when present.
    # Estimated (MCP) rows measure payload size — those bytes are already part of
    # the LLM's context window, so summing them with the LLM total would double-count.
    runner_indices = [i for i, r in enumerate(records) if r.get("server") == "agent-runner"]
    total_rows     = [records[i] for i in runner_indices] if runner_indices else records

    for r in total_rows:
        total_in   += r.get("input_tokens", 0)
        total_out  += r.get("output_tokens", 0)
        total_cost += r.get("cost_usd", 0.0)

    all_exact = all(r.get("token_method") == "exact" for r in total_rows)
    in_tot    = _fmt_tok(total_in, all_exact)
    out_tot   = _fmt_tok(total_out, all_exact)
    pfx_tot   = "" if all_exact else "~"
    cost_tot  = f"{pfx_tot}${total_cost:.6f}"

    if runner_indices and len(runner_indices) < len(records):
        row_nums  = [i + 1 for i in runner_indices]
        est_count = len(records) - len(runner_indices)
        row_label = f"row {row_nums[0]}" if len(row_nums) == 1 else f"rows {','.join(map(str, row_nums))}"
        tot_label = f"Session Total ({row_label}; {est_count} MCP rows in LLM ctx)"
    else:
        tot_label = "Session Total"

    p = get_pricing(m)
    rate_note = (
        f"  Model: {m} — rates applied: "
        f"${p.input_per_mtok:.2f}/M input, ${p.output_per_mtok:.2f}/M output."
        if p else f"  Model: {m}"
    )
    if runner_indices and len(runner_indices) < len(records):
        rate_note += (
            "\n  Session total = exact LLM API cost only."
            " Estimated MCP rows are payload-size heuristics already included in the LLM context."
        )
    elif not all_exact:
        rate_note += (
            "\n  Estimates (~) based on payload size; "
            "exact rows come directly from the LLM API response."
        )

    print(f"  {'─'*(W-2)}", file=sys.stderr)
    print(
        f"  {'':>{S1}}  {c['bold']}{tot_label:<{S2}}{c['reset']}  "
        f"{'':>{S3}}  {c['bold']}{in_tot:>{S4}}  {out_tot:>{S5}}  "
        f"{c['green']}{cost_tot:>12}{c['reset']}",
        file=sys.stderr,
    )
    print(f"{c['dim']}\n{rate_note}{c['reset']}", file=sys.stderr)
    print(f"{c['bold']}{'═'*W}{c['reset']}\n", file=sys.stderr)


# ── Session summary on exit (MCP server / Claude Code path) ───────────────

def _print_summary() -> None:
    # Suppressed when agent_runner will print the unified pipeline summary
    if os.environ.get("STLC_COST_SUMMARY", "").lower() == "suppress":
        return
    if not _TRACKING_ENABLED or _session is None or not _session.records:
        return

    sess    = _session
    records = sess.records
    elapsed = time.time() - sess.started_at
    c       = _C

    total_cost   = sum(r.get("cost_usd", 0.0) for r in records)
    total_tokens = sum(r.get("estimated_tokens", 0) for r in records)

    W = 68
    print(f"\n{c['bold']}{'═'*W}{c['reset']}", file=sys.stderr)
    print(f"{c['bold']}  stlc-agents · Cost Summary  ·  {sess.id}{c['reset']}", file=sys.stderr)
    print(f"{c['bold']}{'═'*W}{c['reset']}", file=sys.stderr)

    print(f"\n  {'Agent':<28} {'Artifact':<34} {'~Tokens':>8}  {'Cost':>10}  {'ms':>6}", file=sys.stderr)
    print(f"  {'─'*W}", file=sys.stderr)
    for r in records:
        raw    = r.get("estimated_tokens", 0)
        tok    = f"{raw/1000:.1f}K" if raw >= 1000 else str(raw)
        pfx    = "" if r.get("token_method") == "exact" else "~"
        art    = _TOOL_ARTIFACT.get(r.get("tool", ""), r.get("tool", "?"))
        print(
            f"  {r.get('server','?'):<28} {art:<34} "
            f"{pfx}{tok:>8}  ${r.get('cost_usd',0):.6f}  {r.get('latency_ms',0):>6}",
            file=sys.stderr,
        )

    tok_total = f"{total_tokens/1000:.1f}K" if total_tokens >= 1000 else str(total_tokens)
    print(f"  {'─'*W}", file=sys.stderr)
    print(
        f"  {c['bold']}{'TOTAL':<28} {'':<34} {tok_total:>9}  "
        f"{c['green']}${total_cost:.6f}{c['reset']}  {elapsed:.1f}s",
        file=sys.stderr,
    )

    model_str = _MODEL_ID
    if _PRICING:
        model_str += f"  (${_PRICING.input_per_mtok}/${_PRICING.output_per_mtok} per MTok in/out)"
    print(f"\n  {c['dim']}Model: {model_str}{c['reset']}", file=sys.stderr)
    print(f"  {c['dim']}Log: {sess.log_path}{c['reset']}", file=sys.stderr)
    print(f"{c['bold']}{'═'*W}{c['reset']}\n", file=sys.stderr)


atexit.register(_print_summary)