"""LLM-based training error analysis and solution suggestions."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

# Common training errors with rule-based solutions
KNOWN_ERRORS: dict[str, dict] = {
    "OutOfMemoryError": {
        "category": "VRAM",
        "suggestions": [
            "Reduce batch_size (try halving it)",
            "Enable gradient checkpointing (use_gradient_checkpointing=True)",
            "Use 4-bit quantization instead of 8-bit or none",
            "Reduce max_seq_length",
            "Reduce LoRA rank",
        ],
    },
    "CUDA out of memory": {
        "category": "VRAM",
        "suggestions": [
            "Reduce batch_size (try halving it)",
            "Enable gradient checkpointing",
            "Use 4-bit quantization",
            "Reduce max_seq_length to match your data's P90",
        ],
    },
    "NCCL error": {
        "category": "distributed",
        "suggestions": [
            "Check that all GPUs are visible (nvidia-smi)",
            "Set NCCL_DEBUG=INFO for detailed error output",
            "Try NCCL_P2P_DISABLE=1 if using different GPU architectures",
            "Ensure torch.distributed is properly initialized",
        ],
    },
    "RuntimeError: Expected all tensors to be on the same device": {
        "category": "device",
        "suggestions": [
            "Ensure model and data are on the same device",
            "Check that LoRA adapters are on the correct device",
            "Use device_map='auto' when loading the model",
        ],
    },
    "ValueError: Tokenizer": {
        "category": "tokenizer",
        "suggestions": [
            "Ensure the tokenizer matches the model",
            "Set pad_token = eos_token if pad_token is not defined",
            "Check that the dataset format matches the expected template",
        ],
    },
    "ImportError": {
        "category": "dependency",
        "suggestions": [
            "Install missing package with pip",
            "Check that the correct version is installed",
            "Try: pip install bilgeai[ml] for all ML dependencies",
        ],
    },
    "KeyError": {
        "category": "data",
        "suggestions": [
            "Check dataset column names match the expected format",
            "Verify dataset format (alpaca needs 'instruction'+'output', "
            "sharegpt needs 'conversations')",
            "Run 'bilge analyze' to inspect your dataset structure",
        ],
    },
    "loss is nan": {
        "category": "training",
        "suggestions": [
            "Reduce learning rate (try 1e-5 or lower)",
            "Use fp32 instead of fp16/bf16 for initial debugging",
            "Check for empty samples in your dataset",
            "Enable gradient clipping (max_grad_norm=1.0)",
        ],
    },
    "loss is 0.0": {
        "category": "training",
        "suggestions": [
            "Check that labels are correctly set (not all -100)",
            "Verify the dataset formatting template",
            "Ensure the model is not in eval mode during training",
        ],
    },
    "DeepSpeed": {
        "category": "deepspeed",
        "suggestions": [
            "Ensure DeepSpeed is installed: pip install deepspeed",
            "Check ZeRO stage compatibility with your setup",
            "Try disabling CPU offloading if you have enough VRAM",
            "Set CUDA_VISIBLE_DEVICES correctly",
        ],
    },
}

ANALYZE_ERROR_PROMPT = """You are an expert ML engineer specialized in LLM fine-tuning debugging.

Analyze this training error and provide actionable solutions.

Error message:
```
{error_message}
```

Error traceback (if available):
```
{traceback}
```

Training configuration:
{config_json}

Hardware:
{hardware_json}

Provide your analysis as a JSON object:
{{
  "category": "<error category: VRAM|distributed|device|tokenizer|dependency|data|training|other>",
  "root_cause": "<1-2 sentence root cause explanation>",
  "suggestions": ["<solution 1>", "<solution 2>", "<solution 3>"],
  "config_changes": {{"<param_name>": "<new_value>"}}
}}

Be specific and actionable. If you suggest config changes, provide exact values."""


@dataclass
class ErrorAnalysis:
    """Result of error analysis."""

    category: str
    root_cause: str
    suggestions: list[str] = field(default_factory=list)
    config_changes: dict = field(default_factory=dict)
    source: str = "rule_engine"


def analyze_error_rule_based(
    error_message: str,
    traceback_str: str = "",
) -> ErrorAnalysis:
    """Analyze a training error using pattern matching (no LLM needed)."""
    full_text = f"{error_message}\n{traceback_str}".lower()

    for pattern, info in KNOWN_ERRORS.items():
        if pattern.lower() in full_text:
            return ErrorAnalysis(
                category=info["category"],
                root_cause=f"Detected known error pattern: {pattern}",
                suggestions=info["suggestions"],
                source="rule_engine",
            )

    return ErrorAnalysis(
        category="unknown",
        root_cause="Could not automatically identify the error pattern.",
        suggestions=[
            "Check the full traceback for more details",
            "Search for the error message in HuggingFace forums",
            "Try running with --dry-run first to validate config",
            "Run 'bilge doctor' to check your environment",
        ],
        source="rule_engine",
    )


def analyze_error_with_llm(
    error_message: str,
    traceback_str: str = "",
    config_dict: dict | None = None,
    hardware_dict: dict | None = None,
    llm_config: dict | None = None,
) -> ErrorAnalysis:
    """Analyze a training error using LLM, with rule-based fallback."""
    rule_result = analyze_error_rule_based(error_message, traceback_str)

    llm_config = llm_config or {}
    try:
        from bilgeai.llm import get_llm_client

        client = get_llm_client(
            provider=llm_config.get("provider", "anthropic"),
            model=llm_config.get("model", "claude-sonnet-4-20250514"),
            api_key_env=llm_config.get("api_key_env", "ANTHROPIC_API_KEY"),
            base_url=llm_config.get("base_url"),
        )

        prompt = ANALYZE_ERROR_PROMPT.format(
            error_message=error_message[:2000],
            traceback=traceback_str[:3000] if traceback_str else "Not available",
            config_json=json.dumps(config_dict or {}, indent=2),
            hardware_json=json.dumps(
                {k: v for k, v in (hardware_dict or {}).items() if k != "gpus"},
                indent=2,
            ),
        )

        response = client.generate(prompt)
        parsed = _parse_analysis(response.content)

        if parsed:
            return ErrorAnalysis(
                category=parsed.get("category", rule_result.category),
                root_cause=parsed.get("root_cause", rule_result.root_cause),
                suggestions=parsed.get("suggestions", rule_result.suggestions),
                config_changes=parsed.get("config_changes", {}),
                source=client.source_name,
            )

    except Exception:
        logger.debug("LLM error analysis failed, using rule engine", exc_info=True)

    return rule_result


def _parse_analysis(content: str) -> dict | None:
    """Parse LLM response to extract error analysis JSON."""
    content = content.strip()

    if "```" in content:
        parts = content.split("```")
        for part in parts:
            part = part.strip()
            if part.startswith("json"):
                part = part[4:].strip()
            try:
                return json.loads(part)
            except (json.JSONDecodeError, ValueError):
                continue

    try:
        return json.loads(content)
    except (json.JSONDecodeError, ValueError):
        return None
