from __future__ import annotations

import ast
import json
import logging
import os
import re
from pathlib import Path
from typing import Any

from amplitude_ai.mcp.contract import MCP_PROMPTS, MCP_RESOURCES, MCP_SERVER_NAME, MCP_TOOLS
from amplitude_ai.mcp.generate_verify_test import generate_verify_test as _generate_verify_test_impl
from amplitude_ai.mcp.instrument_file import instrument_file as _instrument_file_impl
from amplitude_ai.mcp.patterns import get_integration_patterns
from amplitude_ai.mcp.scan_project import scan_project as _scan_project

logger = logging.getLogger(__name__)

_catalog_cache: dict[str, Any] | None = None


def _read_event_catalog() -> dict[str, Any]:
    global _catalog_cache
    if _catalog_cache is not None:
        return _catalog_cache
    catalog_path = Path(__file__).resolve().parent.parent / "data" / "agent_event_catalog.json"
    with catalog_path.open(encoding="utf-8") as file:
        _catalog_cache = json.load(file)
    return _catalog_cache  # type: ignore[return-value]


def _get_event_schema(event_type: str | None = None) -> dict[str, Any]:
    catalog = _read_event_catalog()
    events = catalog.get("events", [])
    if not isinstance(events, list):
        events = []
    if event_type:
        events = [event for event in events if isinstance(event, dict) and event.get("event_type") == event_type]
    return {"total": len(events), "events": events}


def _get_integration_pattern(pattern_id: str | None = None) -> dict[str, Any]:
    patterns = get_integration_patterns()
    if pattern_id:
        patterns = [pattern for pattern in patterns if pattern.get("id") == pattern_id]
    return {"patterns": patterns}


def _validate_setup() -> dict[str, Any]:
    required_env = ["AMPLITUDE_AI_API_KEY"]
    missing_env = [name for name in required_env if not os.environ.get(name)]
    return {
        "status": "ok" if not missing_env else "missing_env",
        "missing_env": missing_env,
    }


def _normalize_content_tier(value: str) -> str:
    if value in {"full", "metadata_only", "customer_enriched"}:
        return value
    return "full"


def _suggest_instrumentation(
    framework: str = "python",
    provider: str = "openai",
    content_tier: str = "full",
) -> dict[str, Any]:
    normalized_tier = _normalize_content_tier(content_tier)
    provider_step = (
        f"use {provider} wrapper/swap integration (for example: `{ 'OpenAI' if provider == 'openai' else 'Anthropic' }(amplitude=ai, ...)`)"
        if provider in ("openai", "anthropic")
        else f"use the {provider} provider wrapper class from amplitude_ai"
    )
    framework_step = (
        "wire AmplitudeAIMiddleware for request identity propagation and execute LLM calls in session context"
        if framework in ("fastapi", "starlette", "flask", "django")
        else "attach BoundAgent + Session where request or conversation identity exists"
    )
    tier_line = (
        "Content tier (`full`): maximum insight and automatic server enrichments. Prefer `redact_pii=True` (+ optional custom patterns)."
        if normalized_tier == "full"
        else (
            "Content tier (`metadata_only`): no content leaves your infrastructure; retain token/cost/latency/session analytics."
            if normalized_tier == "metadata_only"
            else "Content tier (`customer_enriched`): no content leaves infra; send your own labels with `track_session_enrichment()`."
        )
    )
    next_step = (
        "if you need quality/topic analytics without sending content, move to `customer_enriched` and emit structured enrichments."
        if normalized_tier == "metadata_only"
        else (
            "if policy allows and you want zero eval-code overhead, consider `full` with redaction for automatic server enrichments."
            if normalized_tier == "customer_enriched"
            else "if policy prohibits content egress, switch to `metadata_only` or `customer_enriched` without changing session instrumentation."
        )
    )

    now = [
        "install `amplitude-ai`",
        "initialize `AmplitudeAI` with API key",
        provider_step,
        "bind session context with `agent = ai.agent(...); with agent.session(...): ...`",
        framework_step,
    ]
    next_items = [
        next_step,
        "treat `patch(amplitude=...)` as migration quickstart, not steady-state production.",
    ]
    why = [
        "session lineage unlocks scoring, enrichments, and reliable product-to-AI funnels.",
        "wrapper/swap integration keeps implementation effort low while preserving field fidelity.",
        tier_line,
        "privacy controls apply before events leave your process.",
    ]
    steps = [
        "Now:",
        *[f"- {item}" for item in now],
        "Next:",
        *[f"- {item}" for item in next_items],
        "Why:",
        *[f"- {item}" for item in why],
        "Validate: run `amplitude-ai-doctor` and verify [Agent] session-scoped events.",
    ]

    return {
        "framework": framework,
        "provider": provider,
        "content_tier": normalized_tier,
        "now": now,
        "next": next_items,
        "why": why,
        "steps": steps,
    }


_LLM_METHOD_CHAINS: dict[tuple[str, ...], tuple[str, str]] = {
    ("chat", "completions", "create"): ("openai", "chat.completions.create"),
    ("chat", "completions", "parse"): ("openai", "chat.completions.parse"),
    ("responses", "create"): ("openai", "responses.create"),
    ("messages", "create"): ("anthropic", "messages.create"),
    ("generate_content",): ("gemini", "generate_content"),
    ("invoke_model",): ("bedrock", "invoke_model"),
    ("converse",): ("bedrock", "converse"),
    ("chat", "complete"): ("mistral", "chat.complete"),
}

_SINGLE_METHOD_PROVIDER_HINTS: dict[str, tuple[str, ...]] = {
    "gemini": ("gemini", "google", "generative", "genai"),
    "bedrock": ("bedrock", "boto3", "botocore", "runtime"),
}


class _LLMCallVisitor(ast.NodeVisitor):
    """AST visitor that finds LLM provider method calls."""

    def __init__(self, has_patch: bool, wrapped_clients: set[str], provider_signals: set[str]) -> None:
        self.call_sites: list[dict[str, object]] = []
        self._has_patch = has_patch
        self._wrapped_clients = wrapped_clients
        self._provider_signals = provider_signals

    def _extract_chain_and_receiver(self, node: ast.expr) -> tuple[list[str], str | None]:
        parts: list[str] = []
        while isinstance(node, ast.Attribute):
            parts.append(node.attr)
            node = node.value
        parts.reverse()
        receiver = node.id if isinstance(node, ast.Name) else None
        return parts, receiver

    def _single_method_has_provider_signal(
        self,
        provider: str,
        chain: list[str],
        receiver: str | None,
    ) -> bool:
        if receiver is not None and receiver in self._wrapped_clients:
            return True
        if provider in self._provider_signals:
            return True
        hints = _SINGLE_METHOD_PROVIDER_HINTS.get(provider, ())
        chain_prefix = chain[:-1]
        for token in chain_prefix:
            lowered = token.lower()
            if any(hint in lowered for hint in hints):
                return True
        return False

    def visit_Call(self, node: ast.Call) -> None:
        if isinstance(node.func, ast.Attribute):
            chain, receiver = self._extract_chain_and_receiver(node.func)
            for method_chain, (provider, api) in _LLM_METHOD_CHAINS.items():
                if len(chain) >= len(method_chain) and tuple(chain[-len(method_chain):]) == method_chain:
                    if len(method_chain) == 1 and not self._single_method_has_provider_signal(provider, chain, receiver):
                        continue
                    instrumented = self._has_patch or (receiver is not None and receiver in self._wrapped_clients)
                    self.call_sites.append({
                        "line": node.lineno,
                        "provider": provider,
                        "api": api,
                        "instrumented": instrumented,
                    })
                    break
        self.generic_visit(node)


def _validate_file(source: str = "", language: str = "python") -> dict[str, Any]:
    """Analyze source code to detect LLM call sites and report instrumentation status.

    Uses Python's ast module for accurate detection when language is 'python',
    falls back to regex for other languages.
    """
    has_patch = bool(re.search(r'\bpatch\s*\(\s*amplitude\s*=', source, re.DOTALL))
    has_session_context = bool(
        re.search(r"\.session\s*\(", source)
        or re.search(r"\bsession\.run\s*\(", source)
        or "AmplitudeAIMiddleware" in source
    )

    wrapped_client_re = re.compile(
        r'(\w+)\s*=\s*(?:OpenAI|Anthropic|Gemini|AzureOpenAI|Bedrock|Mistral)\s*\(.*?amplitude\s*='
    )
    wrapped_clients = {m.group(1) for m in wrapped_client_re.finditer(source)}
    wrap_re = re.compile(r'(\w+)\s*=\s*(?:\w+\.)*wrap\s*\(.*?amplitude\s*=')
    wrapped_clients |= {m.group(1) for m in wrap_re.finditer(source)}

    has_amplitude_import = bool(
        re.search(r'from\s+amplitude_ai\s+import', source)
        or re.search(r'import\s+amplitude_ai', source)
        or re.search(r'\bAmplitudeAI\b', source)
    )
    provider_signals: set[str] = set()
    if re.search(r'\bopenai\b|from\s+openai\s+import|import\s+openai', source, re.IGNORECASE):
        provider_signals.add("openai")
    if re.search(r'\banthropic\b|from\s+anthropic\s+import|import\s+anthropic', source, re.IGNORECASE):
        provider_signals.add("anthropic")
    if re.search(r'google\.generativeai|google\.genai|\bGemini\b', source, re.IGNORECASE):
        provider_signals.add("gemini")
    if re.search(r'\bboto3\b|\bbotocore\b|\bBedrock\b', source, re.IGNORECASE):
        provider_signals.add("bedrock")
    if re.search(r'\bmistral\b|\bmistralai\b|\bMistral\b', source, re.IGNORECASE):
        provider_signals.add("mistral")

    call_sites: list[dict[str, object]] = []

    _PROVIDER_CONSTRUCTOR_NAMES = {"OpenAI", "Anthropic", "Gemini", "AzureOpenAI", "Bedrock", "Mistral"}

    if language == "python":
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
                    cfunc = node.value.func
                    is_patch_call = (isinstance(cfunc, ast.Name) and cfunc.id == "patch") or (
                        isinstance(cfunc, ast.Attribute) and cfunc.attr == "patch"
                    )
                    if is_patch_call and any(kw.arg == "amplitude" for kw in node.value.keywords):
                        has_patch = True
                if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
                    func = node.value.func
                    is_wrap = (isinstance(func, ast.Name) and func.id == "wrap") or (
                        isinstance(func, ast.Attribute) and func.attr == "wrap"
                    )
                    if is_wrap and any(kw.arg == "amplitude" for kw in node.value.keywords):
                        for target in node.targets:
                            if isinstance(target, ast.Name):
                                wrapped_clients.add(target.id)
                        continue

                    func_name = None
                    if isinstance(func, ast.Name):
                        func_name = func.id
                    elif isinstance(func, ast.Attribute):
                        func_name = func.attr
                    if func_name in _PROVIDER_CONSTRUCTOR_NAMES:
                        has_amplitude_kwarg = any(
                            kw.arg == "amplitude" for kw in node.value.keywords
                        )
                        if has_amplitude_kwarg:
                            for target in node.targets:
                                if isinstance(target, ast.Name):
                                    wrapped_clients.add(target.id)
            visitor = _LLMCallVisitor(has_patch, wrapped_clients, provider_signals)
            visitor.visit(tree)
            call_sites = visitor.call_sites
        except SyntaxError:
            call_sites = _validate_file_regex(source, has_patch, wrapped_clients)
    else:
        call_sites = _validate_file_regex(source, has_patch, wrapped_clients)

    uninstrumented = [s for s in call_sites if not s["instrumented"]]
    if uninstrumented:
        suggestions = [
            "Now: instrument provider calls with wrapper/swap import (for example, `OpenAI(amplitude=ai, api_key=...)`).",
            "Next: add session lineage with `agent = ai.agent(...); with agent.session(...): ...`.",
            "Why: session context unlocks scoring, enrichments, and reliable session funnels.",
            "Content tiers: choose `content_mode` (`full`, `metadata_only`, or `customer_enriched`) and prefer `redact_pii=True` with `full`.",
            "Fallback: use `patch(amplitude=...)` for migration speed, then graduate to wrapper + session context.",
        ]
    elif call_sites and not has_session_context:
        suggestions = [
            "Tracking is present, but session context is missing.",
            "Now: execute LLM calls within `agent.session(...)` (or middleware-backed request context).",
            "Why: without session lineage you lose high-value outcomes like session enrichments, scoring, and reliable funnels.",
            "Content tiers: set `content_mode` intentionally and keep `redact_pii=True` when using `full`.",
        ]
    elif call_sites:
        suggestions = [
            "File appears instrumented with session lineage.",
            "For privacy-by-default, set `content_mode` intentionally and enable `redact_pii=True` when using `full`.",
        ]
    else:
        suggestions = [
            "No supported LLM call sites detected in this file.",
            "If this file should emit AI telemetry, add wrapped provider calls and session context.",
        ]

    return {
        "total_call_sites": len(call_sites),
        "instrumented": len(call_sites) - len(uninstrumented),
        "uninstrumented": len(uninstrumented),
        "has_amplitude_import": has_amplitude_import,
        "has_session_context": has_session_context,
        "call_sites": call_sites,
        "suggestions": suggestions,
    }


def _validate_file_regex(source: str, has_patch: bool, wrapped_clients: set[str]) -> list[dict[str, object]]:
    """Regex fallback for non-Python source or when AST parsing fails."""
    llm_patterns = [
        (r"\.chat\.completions\.create\s*\(", r"(\w+)(?:\.\w+)*\.chat\.completions\.create\s*\(", "openai", "chat.completions.create"),
        (r"\.chat\.completions\.parse\s*\(", r"(\w+)(?:\.\w+)*\.chat\.completions\.parse\s*\(", "openai", "chat.completions.parse"),
        (r"\.responses\.create\s*\(", r"(\w+)(?:\.\w+)*\.responses\.create\s*\(", "openai", "responses.create"),
        (r"\.messages\.create\s*\(", r"(\w+)(?:\.\w+)*\.messages\.create\s*\(", "anthropic", "messages.create"),
        (r"\.generate_content\s*\(", r"(\w+)(?:\.\w+)*\.generate_content\s*\(", "gemini", "generate_content"),
        (r"\.invoke_model\s*\(", r"(\w+)(?:\.\w+)*\.invoke_model\s*\(", "bedrock", "invoke_model"),
        (r"\.converse\s*\(", r"(\w+)(?:\.\w+)*\.converse\s*\(", "bedrock", "converse"),
        (r"\.chat\.complete\s*\(", r"(\w+)(?:\.\w+)*\.chat\.complete\s*\(", "mistral", "chat.complete"),
    ]
    call_sites: list[dict[str, object]] = []
    for i, line in enumerate(source.split("\n"), start=1):
        for detect_pattern, receiver_pattern, provider, api in llm_patterns:
            if re.search(detect_pattern, line):
                rm = re.search(receiver_pattern, line)
                receiver = rm.group(1) if rm else ""
                instrumented = has_patch or receiver in wrapped_clients
                call_sites.append({"line": i, "provider": provider, "api": api, "instrumented": instrumented})
    return call_sites


def _heading_priority(heading: str) -> int:
    normalized = heading.lower()
    if "choose your integration tier" in normalized or "privacy & content control" in normalized:
        return 3
    if "bound agent" in normalized or "boundagent" in normalized or "session" in normalized:
        return 2
    return 1


def _search_docs(query: str = "", max_results: int = 5) -> dict[str, Any]:
    """Search SDK README and API reference for a keyword or phrase."""
    if not query:
        return {"error": "query is required"}

    query_lower = query.lower()
    package_root = Path(__file__).resolve().parent.parent
    sources = [
        ("README.md", package_root.parent / "README.md"),
        ("llms-full.txt", package_root.parent / "llms-full.txt"),
    ]

    results: list[dict[str, Any]] = []
    for name, path in sources:
        if not path.exists():
            continue
        lines = path.read_text(encoding="utf-8").split("\n")
        current_heading = "(top)"
        for i, line in enumerate(lines):
            if line.startswith("#"):
                current_heading = line.lstrip("#").strip()
            if query_lower in line.lower():
                start = max(0, i - 2)
                end = min(len(lines), i + 3)
                results.append({
                    "source": name,
                    "heading": current_heading,
                    "snippet": "\n".join(lines[start:end]),
                    "line": i + 1,
                    "priority": _heading_priority(current_heading),
                })

    results.sort(key=lambda item: (-int(item.get("priority", 1)), int(item.get("line", 0))))
    selected = results[:max_results]
    normalized_results = [{k: v for k, v in item.items() if k != "priority"} for item in selected]
    return {"query": query, "total": len(normalized_results), "results": normalized_results}


def _register_resources_and_prompts(mcp: Any) -> None:
    resource_decorator = getattr(mcp, "resource", None)
    if callable(resource_decorator):
        try:
            @resource_decorator(MCP_RESOURCES["event_schema"])
            def _event_schema_resource() -> str:
                return json.dumps(_read_event_catalog(), indent=2)
        except TypeError as exc:
            logger.warning("Failed to register event_schema resource: %s", exc)

        try:
            @resource_decorator(MCP_RESOURCES["integration_patterns"])
            def _integration_patterns_resource() -> str:
                return json.dumps({"patterns": get_integration_patterns()}, indent=2)
        except TypeError as exc:
            logger.warning("Failed to register integration_patterns resource: %s", exc)

        try:
            @resource_decorator(MCP_RESOURCES["instrument_guide"])
            def _instrument_guide_resource() -> str:
                guide_path = Path(__file__).resolve().parent.parent.parent / "amplitude-ai.md"
                if guide_path.exists():
                    return guide_path.read_text(encoding="utf-8")
                return "Instrument guide not found. Run `amplitude-ai --print-guide` to see the guide."
        except TypeError as exc:
            logger.warning("Failed to register instrument_guide resource: %s", exc)

    prompt_decorator = getattr(mcp, "prompt", None)
    if callable(prompt_decorator):
        try:
            @prompt_decorator(name=MCP_PROMPTS["instrument_app"])
            def _instrument_app_prompt(framework: str = "python", provider: str = "openai") -> str:
                return (
                    f"Instrument this {framework} app with {provider} using amplitude-ai. "
                    "Start with the production default (wrapper/swap plus ai.agent(...).session(...)), choose a content tier "
                    "(full/metadata_only/customer_enriched) with privacy defaults, and treat patch() only as migration fallback. "
                    "Then identify uninstrumented call sites and run validate_setup."
                )
        except TypeError as exc:
            logger.warning("Failed to register instrument_app prompt: %s", exc)


def _generate_verify_test(scan_result: str) -> str:
    """Generate a pytest verification test file from a scan_project JSON result."""
    parsed: dict[str, Any] = json.loads(scan_result)
    return _generate_verify_test_impl(parsed)


def _instrument_file(
    source: str,
    file_path: str,
    tier: str,
    bootstrap_import: str,
    agent_id: str,
    description: str | None = None,
    providers: list[str] | None = None,
) -> str:
    """Apply instrumentation transforms to a Python source file."""
    return _instrument_file_impl(
        source=source,
        file_path=file_path,
        tier=tier,
        bootstrap_import=bootstrap_import,
        agent_id=agent_id,
        description=description,
        providers=providers,
    )


def build_server() -> Any:
    try:
        from fastmcp import FastMCP
    except ImportError as exc:
        raise RuntimeError(
            "fastmcp is not installed. Install with: pip install 'amplitude-ai[mcp]'"
        ) from exc

    mcp = FastMCP(name=MCP_SERVER_NAME)

    mcp.tool(_get_event_schema, name=MCP_TOOLS["get_event_schema"])
    mcp.tool(_get_integration_pattern, name=MCP_TOOLS["get_integration_pattern"])
    mcp.tool(_validate_setup, name=MCP_TOOLS["validate_setup"])
    mcp.tool(_suggest_instrumentation, name=MCP_TOOLS["suggest_instrumentation"])
    mcp.tool(_validate_file, name=MCP_TOOLS["validate_file"])
    mcp.tool(_search_docs, name=MCP_TOOLS["search_docs"])
    mcp.tool(_scan_project, name=MCP_TOOLS["scan_project"])
    mcp.tool(_generate_verify_test, name=MCP_TOOLS["generate_verify_test"])
    mcp.tool(_instrument_file, name=MCP_TOOLS["instrument_file"])

    _register_resources_and_prompts(mcp)
    return mcp


def run_stdio_server() -> None:
    mcp = build_server()
    mcp.run(transport="stdio")
