"""
LLM client — sends code context to Claude or GPT for semantic analysis.

The LLM's job:
  1. Infer the purpose of each module / package
  2. Identify architectural patterns (MVC, microservices, event-driven, etc.)
  3. Label layers (API, business logic, data access, infrastructure)
  4. Suggest a high-level component breakdown
  5. Generate a Mermaid architecture diagram
"""

from __future__ import annotations

import json
import logging
import os
import time
from dataclasses import dataclass, field

from hld_generator.config import HLDConfig
from hld_generator.graph import CodeGraph, GraphBuilder

logger = logging.getLogger(__name__)

# ── System prompt for the LLM ─────────────────────────────────────────────────

SYSTEM_PROMPT = """\
You are a senior software architect. Given a structured summary of a codebase \
(modules, classes, functions, imports, API endpoints, and a dependency graph), \
produce a **High-Level Design (HLD)** document.

Your output MUST be a JSON object with exactly these keys:

{
  "project_name": "<inferred project name>",
  "summary": "<2-3 sentence overview of what this system does>",
  "architecture_pattern": "<e.g. Monolith, Microservices, MVC, Layered, Event-Driven, Serverless, Hexagonal, etc.>",
  "components": [
    {
      "name": "<component / subsystem name>",
      "purpose": "<what it does in 1-2 sentences>",
      "key_modules": ["<file paths>"],
      "layer": "<one of: API, Business Logic, Data Access, Infrastructure, UI, Shared/Utils, Config, Testing>"
    }
  ],
  "data_flow": [
    {
      "from": "<component name>",
      "to": "<component name>",
      "description": "<what data flows and how>"
    }
  ],
  "external_dependencies": [
    {
      "name": "<library or service>",
      "purpose": "<why it is used>"
    }
  ],
  "tech_stack": {
    "languages": ["<lang>"],
    "frameworks": ["<framework>"],
    "databases": ["<db if detectable>"],
    "infrastructure": ["<docker, k8s, etc. if detectable>"]
  },
  "mermaid_diagram": "<A complete Mermaid flowchart (graph TD) diagram showing components and their relationships. Use subgraphs for layers.>",
  "observations": ["<any notable architectural observations, risks, or suggestions>"]
}

Rules:
- Respond ONLY with valid JSON, no markdown fences.
- Be specific to THIS codebase; do not invent modules or files that were not provided.
- The mermaid_diagram must be syntactically valid Mermaid.
- Group related files into logical components; don't list every file as its own component.
- Identify real architectural layers, not just directories.
"""


@dataclass
class LLMAnalysis:
    """Parsed LLM response."""

    project_name: str = "Unknown Project"
    summary: str = ""
    architecture_pattern: str = ""
    components: list[dict] = field(default_factory=list)
    data_flow: list[dict] = field(default_factory=list)
    external_dependencies: list[dict] = field(default_factory=list)
    tech_stack: dict = field(default_factory=dict)
    mermaid_diagram: str = ""
    observations: list[str] = field(default_factory=list)
    raw_response: str = ""
    error: str = ""


class LLMClient:
    """Send code graph context to an LLM and parse the structured response."""

    # Retry configuration
    MAX_RETRIES = 3
    RETRY_BACKOFF = (1, 2, 4)  # seconds between retries
    TIMEOUT = 60  # seconds

    def __init__(self, config: HLDConfig) -> None:
        self.config = config
        self.provider = config.llm_provider
        self.model = config.llm_model
        self.api_key = config.api_key or self._env_key()
        self._anthropic_client = None
        self._openai_client = None

    def analyse(self, code_graph: CodeGraph) -> LLMAnalysis:
        """Run the LLM analysis pipeline."""
        from hld_generator.fallback import FallbackAnalyzer
        from hld_generator.plugins import registry

        fallback = FallbackAnalyzer(self.config)

        if self.provider == "none":
            return fallback.analyse(code_graph)

        # Check for plugin-registered LLM provider first
        plugin_provider = registry.get_llm_provider(self.provider)

        if not plugin_provider and not self.api_key:
            logger.warning("No API key provided — running without LLM analysis")
            return fallback.analyse(code_graph)

        context = GraphBuilder.graph_summary(code_graph)
        # Add file-level detail (truncated to fit context window)
        context += self._build_module_details(code_graph)

        logger.info("Sending %d chars of context to %s/%s", len(context), self.provider, self.model)

        try:
            if plugin_provider:
                raw = self._call_with_retry(
                    lambda ctx: plugin_provider.call(ctx, SYSTEM_PROMPT), context
                )
            elif self.provider == "anthropic":
                raw = self._call_with_retry(self._call_anthropic, context)
            elif self.provider == "openai":
                raw = self._call_with_retry(self._call_openai, context)
            else:
                logger.warning("Unknown provider '%s' — running without LLM analysis", self.provider)
                return fallback.analyse(code_graph)
        except Exception as exc:
            logger.error("LLM call failed after retries: %s", exc)
            analysis = fallback.analyse(code_graph)
            analysis.error = str(exc)
            return analysis

        return self._parse_response(raw)

    # ── Provider calls ─────────────────────────────────────────────────────

    # Exception types that should NOT be retried (auth / validation errors)
    _NON_RETRYABLE = (
        "AuthenticationError",
        "PermissionDeniedError",
        "BadRequestError",
        "InvalidRequestError",
        "NotFoundError",
    )

    def _call_with_retry(self, call_fn, context: str) -> str:
        """Call an LLM provider with exponential backoff retry.

        Non-retryable errors (authentication, validation) are raised immediately
        to avoid wasting time on requests that will never succeed.
        """
        last_exc = None
        for attempt in range(self.MAX_RETRIES):
            try:
                return call_fn(context)
            except Exception as exc:
                # Fail immediately on auth / validation errors
                exc_type = type(exc).__name__
                if exc_type in self._NON_RETRYABLE:
                    raise
                last_exc = exc
                if attempt < self.MAX_RETRIES - 1:
                    delay = self.RETRY_BACKOFF[min(attempt, len(self.RETRY_BACKOFF) - 1)]
                    logger.warning(
                        "LLM call attempt %d/%d failed: %s. Retrying in %ds...",
                        attempt + 1, self.MAX_RETRIES, exc, delay,
                    )
                    time.sleep(delay)
        raise last_exc  # type: ignore[misc]

    def _call_anthropic(self, context: str) -> str:
        import anthropic

        if self._anthropic_client is None:
            self._anthropic_client = anthropic.Anthropic(api_key=self.api_key, timeout=self.TIMEOUT)
        response = self._anthropic_client.messages.create(
            model=self.model,
            max_tokens=8192,
            system=SYSTEM_PROMPT,
            messages=[{"role": "user", "content": context}],
        )
        return response.content[0].text

    def _call_openai(self, context: str) -> str:
        import openai

        if self._openai_client is None:
            self._openai_client = openai.OpenAI(api_key=self.api_key, timeout=self.TIMEOUT)
        response = self._openai_client.chat.completions.create(
            model=self.model,
            max_tokens=8192,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": context},
            ],
        )
        return response.choices[0].message.content or ""

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

    def _env_key(self) -> str:
        if self.provider == "anthropic":
            return os.environ.get("ANTHROPIC_API_KEY", "")
        elif self.provider == "openai":
            return os.environ.get("OPENAI_API_KEY", "")
        return ""

    def _build_module_details(self, cg: CodeGraph) -> str:
        """Build a truncated detail section for the LLM."""
        lines = ["\n\n## Module Details\n"]
        char_budget = self.config.max_tokens_per_chunk * 3  # rough chars-per-token
        accumulated_len = len(lines[0])
        modules_added = 0

        sorted_modules = sorted(cg.modules.items())
        for mod_id, mod in sorted_modules:
            block = []
            block.append(f"\n### `{mod_id}` ({mod.language}, {mod.line_count} lines)")
            if mod.summary:
                block.append(f"  Summary: {mod.summary[:200]}")
            if mod.classes:
                block.append(f"  Classes: {', '.join(mod.classes)}")
            if mod.functions:
                fns = mod.functions[:20]
                extra = f" (+{len(mod.functions) - 20} more)" if len(mod.functions) > 20 else ""
                block.append(f"  Functions: {', '.join(fns)}{extra}")
            if mod.endpoints:
                block.append(f"  Endpoints: {', '.join(mod.endpoints)}")
            if mod.imports_raw:
                imps = mod.imports_raw[:15]
                block.append(f"  Imports: {', '.join(imps)}")

            section = "\n".join(block)
            if accumulated_len + 1 + len(section) > char_budget:
                remaining = len(sorted_modules) - modules_added
                lines.append(f"\n... ({remaining} modules truncated for context limit)")
                break
            lines.append(section)
            accumulated_len += 1 + len(section)  # +1 for joining newline
            modules_added += 1

        return "\n".join(lines)

    def _parse_response(self, raw: str) -> LLMAnalysis:
        """Parse the JSON response from the LLM."""
        analysis = LLMAnalysis(raw_response=raw)

        # Strip markdown fences if present
        cleaned = raw.strip()
        if cleaned.startswith("```"):
            cleaned = cleaned.split("\n", 1)[1] if "\n" in cleaned else cleaned[3:]
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3]
        cleaned = cleaned.strip()
        if cleaned[:4].lower() == "json":
            cleaned = cleaned[4:].strip()

        try:
            data = json.loads(cleaned)
        except json.JSONDecodeError as exc:
            logger.warning("Failed to parse LLM JSON: %s", exc)
            analysis.error = f"JSON parse error: {exc}"
            analysis.summary = raw[:500]
            return analysis

        analysis.project_name = str(data.get("project_name", "Unknown Project"))
        analysis.summary = str(data.get("summary", ""))
        analysis.architecture_pattern = str(data.get("architecture_pattern", ""))

        raw_components = data.get("components", [])
        analysis.components = raw_components if isinstance(raw_components, list) else []

        raw_data_flow = data.get("data_flow", [])
        analysis.data_flow = raw_data_flow if isinstance(raw_data_flow, list) else []

        raw_ext_deps = data.get("external_dependencies", [])
        analysis.external_dependencies = raw_ext_deps if isinstance(raw_ext_deps, list) else []

        raw_tech = data.get("tech_stack", {})
        analysis.tech_stack = raw_tech if isinstance(raw_tech, dict) else {}

        analysis.mermaid_diagram = str(data.get("mermaid_diagram", ""))

        raw_obs = data.get("observations", [])
        analysis.observations = raw_obs if isinstance(raw_obs, list) else []

        return analysis
