"""Prompt and config analyzer for generating optimization suggestions."""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

import litellm

from sage_evaluator.exceptions import SuggestionError
from sage_evaluator.models import Suggestion, SuggestionCategory, SuggestionReport

if TYPE_CHECKING:
    from sage.main_config import MainConfig

logger = logging.getLogger(__name__)

_ANALYSIS_PROMPT_TEMPLATE = """\
You are an expert Sage agent configuration reviewer. Analyze the following agent
configuration and skills, then produce actionable optimization suggestions.

## Agent Configuration

Name: {name}
Model: {model}
Description: {description}
Max Turns: {max_turns}
Permission: {permission}
Extensions: {extensions}
Skills Dir: {skills_dir}

### System Prompt (markdown body)
{system_prompt}

{skills_section}

## Instructions

Analyze the agent thoroughly and return a JSON object with a "suggestions" array.
Each suggestion must have:
  - "category": one of "prompt_improvement", "tool_extraction", "guardrail", "architecture"
  - "title": short title (under 80 chars)
  - "description": 1-3 sentence explanation of the issue and recommended fix
  - "before": the original text snippet (may be empty if not applicable)
  - "after": the improved version or proposed replacement (may be empty if not applicable)
  - "impact": "high", "medium", or "low"

Category guidance:
  - "prompt_improvement": Rewrite vague, repetitive, or ambiguous instructions.
    Provide concrete before/after rewrites of the system prompt.
  - "tool_extraction": Identify verbose inline procedures inside the system prompt
    that should instead be callable @tool functions (e.g. multi-step workflows,
    lookup tables, or external calls described in prose).
  - "guardrail": Suggest deterministic validation or safety tools that should be
    added to prevent errors, out-of-scope requests, or unsafe outputs.
  - "architecture": Recommend model selection improvements, subagent decomposition,
    memory configuration, or max_turns tuning.

Return ONLY valid JSON — no markdown fences, no commentary.

Example output shape:
{{
  "suggestions": [
    {{
      "category": "prompt_improvement",
      "title": "Clarify response language requirement",
      "description": "The instruction 'be concise' is vague. Specifying a maximum word count removes ambiguity.",
      "before": "Be concise in your responses.",
      "after": "Limit each response to at most 150 words unless more detail is explicitly requested.",
      "impact": "medium"
    }}
  ]
}}
"""

_SKILLS_SECTION_TEMPLATE = """\
### Loaded Skills
{skills_list}
"""


def _build_analysis_prompt(config: Any, skills: list[Any]) -> str:
    """Construct the analysis prompt from config and skills."""
    # Format permission
    if config.permission:
        perm_parts = []
        for category, value in config.permission:
            if value is not None:
                perm_parts.append(f"{category}: {value}")
        permission_str = ", ".join(perm_parts) if perm_parts else "(none)"
    else:
        permission_str = "(none)"

    extensions_str = ", ".join(config.extensions) if config.extensions else "(none)"
    system_prompt = config._body.strip() if config._body else "(empty)"

    if skills:
        skill_entries = "\n".join(
            f"- **{s.name}**: {s.description}\n  Content preview: {s.content[:200]}..."
            if len(s.content) > 200
            else f"- **{s.name}**: {s.description}\n  Content: {s.content}"
            for s in skills
        )
        skills_section = _SKILLS_SECTION_TEMPLATE.format(skills_list=skill_entries)
    else:
        skills_section = ""

    return _ANALYSIS_PROMPT_TEMPLATE.format(
        name=config.name,
        model=config.model,
        description=config.description or "(none)",
        max_turns=config.max_turns,
        permission=permission_str,
        extensions=extensions_str,
        skills_dir=config.skills_dir or "(none)",
        system_prompt=system_prompt,
        skills_section=skills_section,
    )


def _strip_markdown_fences(text: str) -> str:
    """Remove markdown code fences (```json ... ```) if present."""
    stripped = text.strip()
    if stripped.startswith("```"):
        # Remove opening fence (e.g. ```json or ```)
        first_newline = stripped.find("\n")
        if first_newline != -1:
            stripped = stripped[first_newline + 1 :]
        # Remove closing fence
        if stripped.rstrip().endswith("```"):
            stripped = stripped.rstrip()[:-3].rstrip()
    return stripped


def _parse_suggestions(raw_json: str) -> list[Suggestion]:
    """Parse the LLM JSON response into a list of Suggestion objects.

    Raises:
        SuggestionError: If the JSON is malformed or the structure is unexpected.
    """
    cleaned = _strip_markdown_fences(raw_json)
    try:
        data = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise SuggestionError(f"Analyzer returned invalid JSON: {exc}") from exc

    if not isinstance(data, dict) or "suggestions" not in data:
        raise SuggestionError(
            "Analyzer response missing top-level 'suggestions' key. "
            f"Got keys: {list(data.keys()) if isinstance(data, dict) else type(data).__name__}"
        )

    suggestions: list[Suggestion] = []
    for idx, item in enumerate(data["suggestions"]):
        try:
            # Validate category before constructing — provides a clearer error.
            raw_category = item.get("category", "")
            try:
                category = SuggestionCategory(raw_category)
            except ValueError:
                valid = [c.value for c in SuggestionCategory]
                raise SuggestionError(
                    f"Suggestion[{idx}] has unknown category '{raw_category}'. "
                    f"Valid values: {valid}"
                )

            suggestions.append(
                Suggestion(
                    category=category,
                    title=item.get("title", ""),
                    description=item.get("description", ""),
                    before=item.get("before", ""),
                    after=item.get("after", ""),
                    impact=item.get("impact", ""),
                )
            )
        except SuggestionError:
            raise
        except Exception as exc:
            raise SuggestionError(f"Failed to parse suggestion[{idx}]: {exc}") from exc

    return suggestions


class PromptAnalyzer:
    """Single-pass LLM analyzer that produces optimization suggestions for an agent config.

    Example usage::

        analyzer = PromptAnalyzer(model="azure_ai/claude-opus-4-6")
        report = await analyzer.analyze("/path/to/agent/AGENTS.md")
        for suggestion in report.suggestions:
            print(suggestion.title)
    """

    def __init__(self, model: str) -> None:
        self.model = model

    async def analyze(
        self, config_path: str | Path, central: MainConfig | None = None
    ) -> SuggestionReport:
        """Analyze an agent config file and return an optimization report.

        Loads the agent config (and skills if ``skills_dir`` is set), sends a
        single structured prompt to the configured LLM, and parses the JSON
        response into :class:`~sage_evaluator.models.Suggestion` objects.

        Args:
            config_path: Path to the agent config file (``AGENTS.md``) or a
                directory containing one.
            central: Optional main config for applying config.toml defaults.

        Returns:
            A :class:`~sage_evaluator.models.SuggestionReport` with the
            parsed suggestions.

        Raises:
            SuggestionError: If config loading, the LLM call, or response
                parsing fails.
        """
        from sage.config import load_config
        from sage.skills.loader import load_skills_from_directory

        resolved_path = Path(config_path)
        logger.info("Analyzing agent config: %s", resolved_path)

        try:
            config = load_config(resolved_path, central=central)
        except Exception as exc:
            raise SuggestionError(f"Failed to load agent config '{config_path}': {exc}") from exc

        skills: list[Any] = []
        if config.skills_dir:
            skills_path = resolved_path if resolved_path.is_dir() else resolved_path.parent
            abs_skills_dir = skills_path / config.skills_dir
            try:
                skills = load_skills_from_directory(abs_skills_dir)
                logger.debug("Loaded %d skills from %s", len(skills), abs_skills_dir)
            except Exception as exc:
                logger.warning("Could not load skills from '%s': %s", abs_skills_dir, exc)

        prompt = _build_analysis_prompt(config, skills)
        logger.debug("Sending analysis prompt to model '%s'", self.model)

        # Only request JSON mode for models known to support it reliably.
        # Other models rely on the prompt instruction to return raw JSON.
        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.0,
        }
        if litellm.supports_response_schema(self.model, None):
            kwargs["response_format"] = {"type": "json_object"}

        try:
            response = await litellm.acompletion(**kwargs)
        except Exception as exc:
            raise SuggestionError(f"LLM call failed for model '{self.model}': {exc}") from exc

        raw_content: str = response.choices[0].message.content or ""
        logger.debug("Received %d chars from analyzer", len(raw_content))

        suggestions = _parse_suggestions(raw_content)
        logger.info("Analysis complete: %d suggestion(s) found", len(suggestions))

        return SuggestionReport(
            config_path=str(resolved_path),
            suggestions=suggestions,
            generated_tools=[],
            analyzer_model=self.model,
        )
