"""
agent_runner.py — Multi-provider LLM agent loop that drives MCP stdio servers directly.

Replaces `claude --mcp-config … -p -- <prompt>` so the webhook pipeline works with
any provider configured in .env (Anthropic, OpenAI, Copilot, Grok, DeepSeek, Azure
OpenAI, Ollama, LM Studio).

Flow
────
1. Start each MCP server subprocess via stdio
2. Collect all tools from every server
3. Call the configured LLM with those tools and the task prompt
4. Loop: execute tool calls → feed results back → repeat until stop/end_turn
"""
from __future__ import annotations

import json
import logging
import os
import time
from contextlib import AsyncExitStack
from typing import Any

from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

from stlc_agents.shared.cost_tracker import get_session_id, print_pipeline_summary, track_llm_call

logger = logging.getLogger("stlc.agent_runner")

MAX_ITERATIONS = 60
MAX_TOKENS     = 8192

# ──────────────────────────────────────────────────────────────────────────────
# Provider resolution
# ──────────────────────────────────────────────────────────────────────────────

_OPENAI_COMPAT: dict[str, tuple[str | None, str, str]] = {
    # provider       base_url                                  key_env            default_model
    "openai":        (None,                                    "OPENAI_API_KEY",  "gpt-4o-mini"),
    "copilot":       ("https://api.githubcopilot.com",         "GITHUB_TOKEN",    "gpt-4o"),
    "grok":          ("https://api.x.ai/v1",                   "XAI_API_KEY",     "grok-3-mini"),
    "deepseek":      ("https://api.deepseek.com/v1",           "DEEPSEEK_API_KEY","deepseek-chat"),
    "azure-openai":  (None,                                    "AZURE_OPENAI_KEY","gpt-4o-mini"),
    "ollama":        ("http://localhost:11434/v1",              "",                "llama3.2"),
    "lm-studio":     ("http://localhost:1234/v1",               "",                "local-model"),
}


def active_provider() -> str:
    """Return the active LLM provider, auto-detecting from env with key-prefix fallback."""
    p = os.environ.get("LLM_PROVIDER", "").strip().lower()
    if p:
        return p
    p = os.environ.get("AI_HEALING_PROVIDER", "").strip().lower()
    if p:
        return p
    # Auto-detect from API key prefix
    key = os.environ.get("AI_HEALING_API_KEY", "").strip()
    if key.startswith("sk-ant-"):
        return "anthropic"
    if key.startswith("gh_") or key.startswith("ghu_"):
        return "copilot"
    if key.startswith("sk-") or key.startswith("sk-proj-"):
        return "openai"
    return "ollama"


def _model_for(provider: str) -> str:
    """Return the configured model for a provider, falling back to the provider default."""
    m = os.environ.get("LLM_MODEL", "").strip()
    if m:
        return m
    cfg = _OPENAI_COMPAT.get(provider.lower(), (None, "", "gpt-4o-mini"))
    return cfg[2]  # default_model slot


def _normalise_copilot_model(model: str) -> str:
    """
    Normalise a model name for the GitHub Copilot API.
    claude-* models are lowercased (no vendor prefix needed).
    Everything else gets the 'openai/' prefix.
    Already-prefixed names (contain '/') are returned unchanged.
    """
    if not model or "/" in model:
        return model
    if "claude" in model.lower():
        return model.lower()
    return f"openai/{model}"


def _resolve_provider(
    provider: str, model: str, api_key: str, base_url: str
) -> tuple[Any, str, str]:
    """Return (async_client, model_name, family) where family is 'anthropic' or 'openai'."""
    p = (provider or "anthropic").lower().strip()

    if p in ("anthropic", "claude", "claude-code"):
        try:
            from anthropic import AsyncAnthropic
        except ImportError:
            raise RuntimeError(
                "anthropic package missing — run: pip install 'qa-stlc-agents[webhook]'"
            )
        key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        client = AsyncAnthropic(api_key=key)
        return client, model or "claude-haiku-4-5-20251001", "anthropic"

    if p == "azure-openai":
        try:
            from openai import AsyncAzureOpenAI
        except ImportError:
            raise RuntimeError(
                "openai package missing — run: pip install 'qa-stlc-agents[webhook]'"
            )
        endpoint = base_url or os.environ.get("AZURE_OPENAI_ENDPOINT", "")
        key = (api_key or os.environ.get("AZURE_OPENAI_KEY")
               or os.environ.get("AZURE_OPENAI_API_KEY", ""))
        api_ver  = os.environ.get("AZURE_OPENAI_API_VERSION", "2024-02-01")
        dep      = model or os.environ.get("AZURE_OPENAI_DEPLOYMENT", "gpt-4o-mini")
        client   = AsyncAzureOpenAI(api_key=key, azure_endpoint=endpoint, api_version=api_ver)
        return client, dep, "openai"

    # All other OpenAI-compatible providers
    try:
        from openai import AsyncOpenAI
    except ImportError:
        raise RuntimeError(
            "openai package missing — run: pip install 'qa-stlc-agents[webhook]'"
        )
    cfg = _OPENAI_COMPAT.get(p, (None, "", "gpt-4o-mini"))
    default_url, key_env, default_model = cfg

    url = (base_url
           or (os.environ.get("LM_STUDIO_URL")  if p == "lm-studio" else None)
           or (os.environ.get("OLLAMA_HOST")     if p == "ollama"    else None)
           or default_url)
    key = api_key or (os.environ.get(key_env) if key_env else "") or "dummy"
    m   = model or default_model

    if p == "copilot":
        m = _normalise_copilot_model(m)

    kw: dict = {"api_key": key}
    if url:
        kw["base_url"] = url
    return AsyncOpenAI(**kw), m, "openai"


# ──────────────────────────────────────────────────────────────────────────────
# MCP server management
# ──────────────────────────────────────────────────────────────────────────────

async def _start_servers(
    mcp_config: dict, env: dict[str, str], stack: AsyncExitStack
) -> dict[str, tuple[ClientSession, Any]]:
    """Start all MCP servers and return {tool_name: (session, mcp_tool)}."""
    servers_cfg: dict = mcp_config.get("mcpServers", {})
    tool_map: dict[str, tuple[ClientSession, Any]] = {}

    for server_name, cfg in servers_cfg.items():
        command = cfg.get("command", "")
        args    = cfg.get("args", [])
        if not command:
            logger.warning("MCP server '%s' has no command — skipping", server_name)
            continue

        params = StdioServerParameters(command=command, args=args, env=env)
        try:
            read, write = await stack.enter_async_context(stdio_client(params))
            session     = await stack.enter_async_context(ClientSession(read, write))
            await session.initialize()
        except Exception as exc:
            logger.error("Failed to start MCP server '%s': %s", server_name, exc)
            continue

        try:
            result = await session.list_tools()
            for tool in result.tools:
                if tool.name in tool_map:
                    logger.debug(
                        "Tool name collision: '%s' already registered; '%s' wins",
                        tool.name, server_name,
                    )
                tool_map[tool.name] = (session, tool)
        except Exception as exc:
            logger.error("Failed to list tools for '%s': %s", server_name, exc)

    return tool_map



def _extract_artefact(tool_name: str, result_text: str, wi_id: str) -> dict | None:
    """Parse a tool result and return a structured artefact dict for the summary report."""
    try:
        data = json.loads(result_text)
    except (json.JSONDecodeError, ValueError):
        return None
    if not isinstance(data, dict):
        return None
    if data.get("success") is False and tool_name != "inspect_helix_project":
        return None

    if tool_name in ("fetch_work_item_for_gherkin", "fetch_work_item", "fetch_feature_hierarchy"):
        title   = data.get("title") or data.get("work_item_title") or f"WI {wi_id}"
        wi_type = data.get("work_item_type") or data.get("type") or "Work Item"
        state   = data.get("state") or ""
        return {
            "name":     title[:26],
            "type":     wi_type,
            "location": f"WI #{wi_id}",
            "detail":   f"State: {state}" if state else "",
        }

    if tool_name in (
        "generate_and_attach_gherkin", "attach_gherkin_to_feature", "attach_gherkin_to_work_item"
    ):
        gherkin    = data.get("gherkin_content") or data.get("content") or ""
        n_sc       = gherkin.lower().count("scenario:")
        att_id     = data.get("attachment_id") or data.get("attachment_url") or ""
        detail     = f"{n_sc} scenario(s)" if n_sc else ""
        if att_id:
            detail = (detail + f"  id:{att_id}").strip()
        return {
            "name":     "Gherkin Scenarios",
            "type":     "Gherkin / attachment",
            "location": f"WI #{wi_id}",
            "detail":   detail,
        }

    if tool_name in ("create_and_link_test_cases", "create_deduped_test_cases"):
        tcs = data.get("test_cases") or data.get("created_test_cases") or []
        n   = len(tcs) if isinstance(tcs, list) else data.get("count", 0)
        return {
            "name":     "Test Cases",
            "type":     "Test Cases",
            "location": f"WI #{wi_id}",
            "detail":   f"{n} test case(s) created",
        }

    if tool_name == "capture_app_context":
        ctx_map = data.get("context_map") or {}
        n_pages = len(ctx_map) if isinstance(ctx_map, dict) else 0
        src     = data.get("locator_source") or ""
        ck      = data.get("cache_key") or ""
        detail  = f"{n_pages} page(s)" + (f"  source:{src}" if src else "")
        return {
            "name":     "App Context",
            "type":     "Browser Snapshot",
            "location": f"cache:{ck}" if ck else "in-memory",
            "detail":   detail,
        }

    if tool_name == "generate_playwright_code":
        ck    = data.get("cache_key") or ""
        files = data.get("files") or {}
        n     = len(files) if isinstance(files, dict) else 0
        names = ", ".join(list(files.keys())[:3]) if isinstance(files, dict) else ""
        if n > 3:
            names += f" +{n - 3} more"
        return {
            "name":     "Playwright Tests",
            "type":     "TypeScript",
            "location": f"cache:{ck}" if ck else "in-memory",
            "detail":   f"{n} file(s): {names}" if names else f"{n} file(s)",
        }

    if tool_name == "write_helix_files":
        helix_root = data.get("helix_root") or ""
        written    = data.get("written") or data.get("files_written") or []
        n          = len(written) if isinstance(written, list) else 0
        names = ", ".join([
            (p.get("dest") or p.get("file_key") or "").rsplit("/", 1)[-1]
            if isinstance(p, dict) else str(p).rsplit("/", 1)[-1]
            for p in written[:3]
        ]) if isinstance(written, list) else ""
        if n > 3:
            names += f" +{n - 3} more"
        return {
            "name":     "Helix QA Files",
            "type":     "TypeScript / Gherkin",
            "location": helix_root or "Helix disk",
            "detail":   f"{n} file(s): {names}" if names else f"{n} file(s)",
        }

    return None


async def _call_mcp_tool(
    tool_map: dict, tool_name: str, arguments: dict, wi_id: str,
    artefacts: list[dict] | None = None,
) -> str:
    """Call an MCP tool and return its result as a string."""
    if tool_name not in tool_map:
        return f"ERROR: unknown tool '{tool_name}'"
    session, _ = tool_map[tool_name]
    try:
        logger.info("claude[wi=%s][tool] %s(%s)", wi_id, tool_name, _abbrev(arguments))
        result = await session.call_tool(tool_name, arguments)
        text = _extract_text(result)
        logger.info("claude[wi=%s][tool-result] %s → %s", wi_id, tool_name, text[:200])
        if artefacts is not None:
            art = _extract_artefact(tool_name, text, wi_id)
            if art is not None:
                artefacts.append(art)
        return text
    except Exception as exc:
        logger.error("MCP tool '%s' raised %s", tool_name, exc)
        return f"ERROR calling {tool_name}: {exc}"


# ──────────────────────────────────────────────────────────────────────────────
# Anthropic agent loop
# ──────────────────────────────────────────────────────────────────────────────

def _anthropic_tools(tool_map: dict) -> list[dict]:
    return [
        {
            "name":         name,
            "description":  tool.description or "",
            "input_schema": tool.inputSchema or {"type": "object", "properties": {}},
        }
        for name, (_, tool) in tool_map.items()
    ]


async def _run_anthropic(
    client: Any, model: str, prompt: str,
    tool_map: dict, wi_id: str,
    artefacts: list[dict] | None = None,
) -> tuple[bool, str]:
    messages = [{"role": "user", "content": prompt}]
    tools    = _anthropic_tools(tool_map)

    total_input_tokens       = 0
    total_output_tokens      = 0
    total_cache_write_tokens = 0
    total_cache_read_tokens  = 0
    t0 = time.monotonic()

    def _flush(iteration: int) -> None:
        track_llm_call(
            model=model, provider="anthropic",
            input_tokens=total_input_tokens,
            output_tokens=total_output_tokens,
            cache_write_tokens=total_cache_write_tokens,
            cache_read_tokens=total_cache_read_tokens,
            work_item_id=wi_id,
            iterations=iteration + 1,
            latency_ms=int((time.monotonic() - t0) * 1000),
        )

    for iteration in range(MAX_ITERATIONS):
        response = await client.messages.create(
            model=model,
            max_tokens=MAX_TOKENS,
            system=(
                "You are a QA automation agent. Execute the requested pipeline steps "
                "by calling the available tools in sequence. Auto-confirm all actions."
            ),
            messages=messages,
            tools=tools or None,
        )

        if response.usage:
            total_input_tokens       += response.usage.input_tokens or 0
            total_output_tokens      += response.usage.output_tokens or 0
            total_cache_write_tokens += getattr(response.usage, "cache_creation_input_tokens", 0) or 0
            total_cache_read_tokens  += getattr(response.usage, "cache_read_input_tokens", 0) or 0

        logger.info(
            "claude[wi=%s][llm] iteration=%d stop=%s in=%d out=%d cache_write=%d cache_read=%d",
            wi_id, iteration, response.stop_reason,
            total_input_tokens, total_output_tokens,
            total_cache_write_tokens, total_cache_read_tokens,
        )

        messages.append({"role": "assistant", "content": response.content})

        if response.stop_reason in ("end_turn", "stop_sequence"):
            _flush(iteration)
            return True, ""

        if response.stop_reason != "tool_use":
            _flush(iteration)
            return False, f"Unexpected stop reason: {response.stop_reason}"

        tool_results = []
        for block in response.content:
            if block.type != "tool_use":
                continue
            result_text = await _call_mcp_tool(tool_map, block.name, block.input, wi_id, artefacts)
            tool_results.append({
                "type":        "tool_result",
                "tool_use_id": block.id,
                "content":     result_text,
            })

        if not tool_results:
            _flush(iteration)
            return False, "tool_use stop reason but no tool_use blocks"

        messages.append({"role": "user", "content": tool_results})

    _flush(MAX_ITERATIONS - 1)
    return False, f"Max iterations ({MAX_ITERATIONS}) reached"


# ──────────────────────────────────────────────────────────────────────────────
# OpenAI-compatible agent loop
# ──────────────────────────────────────────────────────────────────────────────

def _openai_tools(tool_map: dict) -> list[dict]:
    return [
        {
            "type": "function",
            "function": {
                "name":        name,
                "description": tool.description or "",
                "parameters":  tool.inputSchema or {"type": "object", "properties": {}},
            },
        }
        for name, (_, tool) in tool_map.items()
    ]


async def _run_openai(
    client: Any, model: str, prompt: str,
    tool_map: dict, wi_id: str,
    artefacts: list[dict] | None = None,
) -> tuple[bool, str]:
    messages: list[dict] = [
        {
            "role":    "system",
            "content": (
                "You are a QA automation agent. Execute the requested pipeline steps "
                "by calling the available tools in sequence. Auto-confirm all actions."
            ),
        },
        {"role": "user", "content": prompt},
    ]
    tools = _openai_tools(tool_map) or None

    total_input_tokens      = 0
    total_output_tokens     = 0
    total_cache_read_tokens = 0
    t0 = time.monotonic()

    def _flush(iteration: int) -> None:
        track_llm_call(
            model=model, provider="openai",
            input_tokens=total_input_tokens,
            output_tokens=total_output_tokens,
            cache_read_tokens=total_cache_read_tokens,
            work_item_id=wi_id,
            iterations=iteration + 1,
            latency_ms=int((time.monotonic() - t0) * 1000),
        )

    for iteration in range(MAX_ITERATIONS):
        response = await client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice="auto" if tools else None,
            max_completion_tokens=MAX_TOKENS,
        )
        msg    = response.choices[0].message
        finish = response.choices[0].finish_reason

        if response.usage:
            # prompt_tokens is the TOTAL (cached + uncached). Split them so each
            # bucket is billed at the correct rate via pricing.py cache_read_per_mtok.
            details  = getattr(response.usage, "prompt_tokens_details", None)
            cached   = (getattr(details, "cached_tokens", 0) or 0) if details else 0
            total_input_tokens      += (response.usage.prompt_tokens or 0) - cached
            total_output_tokens     += response.usage.completion_tokens or 0
            total_cache_read_tokens += cached

        logger.info(
            "claude[wi=%s][llm] iteration=%d finish=%s in=%d out=%d cached=%d",
            wi_id, iteration, finish,
            total_input_tokens, total_output_tokens, total_cache_read_tokens,
        )

        messages.append(msg.model_dump(exclude_unset=True))

        if finish == "stop":
            _flush(iteration)
            return True, ""

        if finish != "tool_calls":
            _flush(iteration)
            return False, f"Unexpected finish reason: {finish}"

        if not msg.tool_calls:
            _flush(iteration)
            return False, "finish_reason=tool_calls but no tool_calls"

        for tc in msg.tool_calls:
            try:
                args = json.loads(tc.function.arguments or "{}")
            except json.JSONDecodeError:
                args = {}
            result_text = await _call_mcp_tool(tool_map, tc.function.name, args, wi_id, artefacts)
            messages.append({
                "role":         "tool",
                "tool_call_id": tc.id,
                "content":      result_text,
            })

    _flush(MAX_ITERATIONS - 1)
    return False, f"Max iterations ({MAX_ITERATIONS}) reached"


# ──────────────────────────────────────────────────────────────────────────────
# Public entry point
# ──────────────────────────────────────────────────────────────────────────────

async def run_agent_loop(
    mcp_config: dict,
    prompt: str,
    llm_provider: str,
    llm_model: str,
    llm_api_key: str,
    llm_base_url: str,
    env: dict[str, str],
    work_item_id: str,
) -> tuple[bool, str]:
    """
    Drive the STLC agent pipeline using the configured LLM provider.
    Returns (success, error_message).
    """
    try:
        client, model, family = _resolve_provider(llm_provider, llm_model, llm_api_key, llm_base_url)
    except RuntimeError as exc:
        return False, str(exc)

    logger.info(
        "agent_runner: wi=%s provider=%s model=%s",
        work_item_id, llm_provider or "anthropic", model,
    )

    t0 = time.time()
    async with AsyncExitStack() as stack:
        # Share session ID with MCP subprocesses so all costs land in one JSONL.
        # Suppress per-subprocess atexit summaries; parent prints the unified table.
        merged_env = {**env}
        if "STLC_SESSION_ID" not in merged_env:
            merged_env["STLC_SESSION_ID"] = get_session_id()
        merged_env["STLC_COST_SUMMARY"] = "suppress"

        tool_map = await _start_servers(mcp_config, merged_env, stack)
        if not tool_map:
            return False, "No MCP tools available — all servers failed to start"

        logger.info(
            "agent_runner: wi=%s tools=[%s]",
            work_item_id, ", ".join(tool_map.keys()),
        )

        artefacts: list[dict] = []
        if family == "anthropic":
            success, error = await _run_anthropic(
                client, model, prompt, tool_map, work_item_id, artefacts
            )
        else:
            success, error = await _run_openai(
                client, model, prompt, tool_map, work_item_id, artefacts
            )

    # AsyncExitStack has closed all MCP servers; all JSONL records are written.
    print_pipeline_summary(
        get_session_id(), str(work_item_id), time.time() - t0, model,
        artefacts or None,
    )
    return success, error


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────

def _abbrev(d: dict, limit: int = 120) -> str:
    s = json.dumps(d, default=str)
    return s[:limit] + "…" if len(s) > limit else s


def _extract_text(result: Any) -> str:
    """Extract a plain-text representation from an MCP CallToolResult."""
    if hasattr(result, "content"):
        parts = []
        for block in result.content:
            if hasattr(block, "text"):
                parts.append(block.text)
            else:
                parts.append(json.dumps(block.model_dump() if hasattr(block, "model_dump") else str(block)))
        return "\n".join(parts)
    return str(result)
