"""LLM-powered training config optimizer."""

from __future__ import annotations

import json
import logging

from bilgeai.analyzer.models import DatasetProfile
from bilgeai.core.exceptions import LLMError
from bilgeai.llm.base import LLMClient
from bilgeai.llm.prompts import GENERATE_CONFIG_PROMPT, SYSTEM_PROMPT
from bilgeai.optimizer.models import HardwareConstraints, TrainingConfig
from bilgeai.profiler.models import HardwareProfile

logger = logging.getLogger(__name__)


class LLMOptimizer:
    """LLM-powered training config optimizer with guardrail enforcement."""

    def __init__(self, client: LLMClient) -> None:
        self.client = client

    @classmethod
    def from_config(cls, llm_config: dict) -> LLMOptimizer:
        """Create an LLMOptimizer from llm config dict."""
        from bilgeai.llm.router 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"),
        )
        return cls(client)

    def optimize(
        self,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
        constraints: HardwareConstraints,
        model_name: str = "meta-llama/Llama-3-8B",
        user_overrides: dict | None = None,
    ) -> TrainingConfig:
        """Generate full TrainingConfig via LLM within guardrail bounds."""
        user_overrides = user_overrides or {}

        prompt = GENERATE_CONFIG_PROMPT.format(
            hardware_json=json.dumps(hardware.model_dump(exclude={"gpus"}), indent=2, default=str),
            constraints_json=json.dumps(constraints.model_dump(exclude={"explanations"}), indent=2),
            dataset_json=json.dumps(
                dataset.model_dump(exclude={"quality_issues", "pii_findings"}),
                indent=2,
                default=str,
            ),
            overrides_json=json.dumps(user_overrides, indent=2) if user_overrides else "{}",
            model_name=model_name,
            max_batch_size=constraints.max_batch_size,
            required_quantization=constraints.required_quantization or "any",
            distributed_strategy=constraints.distributed_strategy or "any",
        )

        # Try LLM call with retry
        config = self._call_llm_with_retry(prompt, constraints, model_name)

        # Apply user overrides (highest priority)
        if user_overrides:
            for key, value in user_overrides.items():
                if hasattr(config, key) and value is not None:
                    setattr(config, key, value)
                    config.add_explanation(
                        decision=f"{key} = {value}",
                        reason="User override from bilge.yaml",
                        confidence=1.0,
                        source="user",
                    )

        # Post-validation: clamp to guardrail bounds
        config = self._validate_against_constraints(config, constraints)

        # Preserve guardrail explanations
        for expl in constraints.explanations:
            config.explanations.insert(0, expl)

        return config

    def _call_llm_with_retry(
        self,
        prompt: str,
        constraints: HardwareConstraints,
        model_name: str,
    ) -> TrainingConfig:
        """Call LLM and parse response. Retry once on parse failure."""
        from bilgeai.llm.remote import _extract_json_from_response

        last_error: Exception | None = None

        for attempt in range(2):
            try:
                if hasattr(self.client, "generate_structured"):
                    response = self.client.generate_structured(prompt, system=SYSTEM_PROMPT)
                else:
                    response = self.client.generate(prompt, system=SYSTEM_PROMPT)

                # Parse response
                content = response.content.strip()
                try:
                    data = json.loads(content)
                except (json.JSONDecodeError, ValueError):
                    data = _extract_json_from_response(content)

                # Extract explanations before Pydantic validation
                raw_explanations = data.pop("explanations", [])

                config = TrainingConfig.model_validate(data)

                # Re-add explanations
                for expl in raw_explanations:
                    if isinstance(expl, dict):
                        config.add_explanation(
                            decision=expl.get("decision", ""),
                            reason=expl.get("reason", ""),
                            confidence=float(expl.get("confidence", 0.8)),
                            source="llm",
                        )

                logger.info("LLM generated TrainingConfig successfully")
                return config

            except LLMError:
                raise
            except Exception as e:
                last_error = e
                if attempt == 0:
                    logger.warning("LLM config parse failed: %s. Retrying...", e)
                    prompt += (
                        f"\n\nPrevious attempt failed with error: {e}\nPlease fix and try again."
                    )

        raise LLMError(
            f"LLM optimization failed: {last_error}",
            suggestion="Check your API key and network connection.",
        )

    @staticmethod
    def _validate_against_constraints(
        config: TrainingConfig, constraints: HardwareConstraints
    ) -> TrainingConfig:
        """Post-validation: clamp LLM output to guardrail bounds."""
        # Batch size must not exceed max
        if config.batch_size > constraints.max_batch_size:
            logger.warning(
                "LLM suggested batch_size=%d, clamping to max=%d",
                config.batch_size,
                constraints.max_batch_size,
            )
            config.batch_size = constraints.max_batch_size

        # Quantization must match if required
        if (
            constraints.required_quantization is not None
            and config.quantization != constraints.required_quantization
        ):
            logger.warning(
                "LLM suggested quantization=%s, forcing to %s",
                config.quantization,
                constraints.required_quantization,
            )
            config.quantization = constraints.required_quantization

        # Distributed strategy must match if required
        if (
            constraints.distributed_strategy is not None
            and config.distributed_strategy != constraints.distributed_strategy
        ):
            config.distributed_strategy = constraints.distributed_strategy

        # GPU count
        config.num_gpus = constraints.num_gpus

        # Gradient checkpointing
        if constraints.gradient_checkpointing_required:
            config.use_gradient_checkpointing = True

        return config
