"""Config Optimizer - GuardrailEngine + LLM-powered HybridEngine."""

from __future__ import annotations

import ast
import json
import logging
import re
from pathlib import Path

from bilgeai.analyzer.models import DatasetProfile
from bilgeai.core.exceptions import LLMError
from bilgeai.optimizer.models import HardwareConstraints, TrainingConfig
from bilgeai.optimizer.rules.vram import (
    DistributedStrategyRule,
    VRAMBatchSizeRule,
    VRAMQuantizationRule,
)
from bilgeai.profiler.models import HardwareProfile

logger = logging.getLogger(__name__)


def _fix_load_dataset_file_type(script: str, file_type: str, is_directory: bool = False) -> str:
    """Fix load_dataset() calls to use the correct file type.

    LLMs sometimes ignore instructions and hardcode 'json' regardless of actual file type.
    This post-processes the script to ensure correctness.
    Also fixes data_files vs data_dir for directory paths.
    """
    fixed = script

    # Fix file type in load_dataset first argument
    fixed = re.sub(
        r'load_dataset\(\s*["\'](?:json|csv|parquet)["\']',
        f'load_dataset("{file_type}"',
        fixed,
    )

    # Fix data_files -> data_dir for directories (and vice versa)
    if is_directory:
        fixed = re.sub(r"\bdata_files\s*=", "data_dir=", fixed)
    else:
        fixed = re.sub(r"\bdata_dir\s*=", "data_files=", fixed)

    if fixed != script:
        logger.info(
            "Post-fixed load_dataset: file_type='%s', is_directory=%s",
            file_type,
            is_directory,
        )
    return fixed


def _detect_file_type(dataset_path: str) -> str:
    """Detect file type from dataset path extension.

    For directories, peeks at the first file inside to determine type.
    """
    p = Path(dataset_path)
    suffix = p.suffix.lower()

    # No extension — might be a directory
    if not suffix and p.is_dir():
        supported = (".parquet", ".jsonl", ".json", ".csv")
        for f in sorted(p.iterdir()):
            if f.is_file() and f.suffix.lower() in supported:
                suffix = f.suffix.lower()
                break

    if suffix == ".parquet":
        return "parquet"
    elif suffix == ".csv":
        return "csv"
    else:
        return "json"


def _build_torchrun_command(config: TrainingConfig) -> str | None:
    """Build torchrun command for multi-GPU/node setups. None for single GPU."""
    if config.num_gpus <= 1 and config.num_nodes <= 1:
        return None
    parts = ["torchrun", f"--nproc_per_node={config.num_gpus}"]
    if config.num_nodes > 1:
        parts.extend(
            [
                f"--nnodes={config.num_nodes}",
                "--rdzv_backend=c10d",
                "--rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT",
            ]
        )
    parts.append("train.py")
    return " ".join(parts)


class GuardrailEngine:
    """Hardware safety guardrails only. Produces constraints for LLM."""

    def __init__(self) -> None:
        self.rules = [
            VRAMQuantizationRule(),
            VRAMBatchSizeRule(),
            DistributedStrategyRule(),
        ]

    def compute_constraints(
        self,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
        model_name: str = "meta-llama/Llama-3-8B",
    ) -> HardwareConstraints:
        """Run guardrail rules, return constraints for LLM context."""
        config = TrainingConfig(model_name=model_name)

        for rule in self.rules:
            config = rule.apply(config, hardware, dataset)

        return HardwareConstraints(
            max_batch_size=config.batch_size,
            required_quantization=config.quantization,
            max_seq_length_safe=config.max_seq_length,
            distributed_strategy=config.distributed_strategy,
            num_gpus=config.num_gpus,
            gradient_checkpointing_required=config.use_gradient_checkpointing,
            explanations=config.explanations,
        )


class HybridEngine:
    """LLM-first optimization engine with hardware guardrails."""

    def __init__(self) -> None:
        self.guardrail_engine = GuardrailEngine()

    def optimize(
        self,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
        model_name: str = "meta-llama/Llama-3-8B",
        user_overrides: dict | None = None,
        llm_config: dict | None = None,
    ) -> TrainingConfig:
        """Run optimization: guardrails -> LLM -> post-validation."""
        llm_config = llm_config or {}

        # Step 1: Compute hardware constraints
        constraints = self.guardrail_engine.compute_constraints(hardware, dataset, model_name)

        # Step 2: LLM optimization
        from bilgeai.optimizer.llm_optimizer import LLMOptimizer

        optimizer = LLMOptimizer.from_config(llm_config)
        config = optimizer.optimize(
            hardware=hardware,
            dataset=dataset,
            constraints=constraints,
            model_name=model_name,
            user_overrides=user_overrides,
        )

        # Compute torchrun command for multi-GPU/node
        config.torchrun_command = _build_torchrun_command(config)

        return config

    def generate_script(
        self,
        config: TrainingConfig,
        hardware: HardwareProfile,
        dataset: DatasetProfile | None = None,
        dataset_path: str = "./train.jsonl",
        dataset_format: str = "alpaca",
        llm_config: dict | None = None,
    ) -> str:
        """Generate training script: template -> LLM review."""
        from bilgeai.generator import CodeGenerator

        generator = CodeGenerator()
        hw_summary = f"{hardware.cpu_name}, {hardware.ram_total_gb:.0f}GB RAM"
        if hardware.gpus:
            gpu = hardware.gpus[0]
            hw_summary += f", {gpu.name} ({gpu.vram_gb:.0f}GB)"
            if hardware.num_gpus > 1:
                hw_summary += f" x{hardware.num_gpus}"

        # Step 1: Generate baseline from template
        script = generator.generate_sft(
            config,
            dataset_path=dataset_path,
            dataset_format=dataset_format,
            hardware_summary=hw_summary,
        )

        # Detect file type and whether path is a directory
        file_type = _detect_file_type(dataset_path)
        is_directory = Path(dataset_path).is_dir() or not Path(dataset_path).suffix

        # Step 2: LLM review (mandatory)
        ds_summary = self._build_dataset_summary(dataset, dataset_path, file_type)
        script = self._review_script_with_llm(
            script, config, hw_summary, ds_summary, llm_config, file_type, is_directory
        )

        return script

    @staticmethod
    def _build_dataset_summary(
        dataset: DatasetProfile | None, dataset_path: str, file_type: str = "json"
    ) -> str:
        """Build a rich dataset summary string for LLM prompts."""
        if dataset is None:
            return f"Path: {dataset_path} | File type: {file_type}"

        parts = [f"Path: {dataset_path}"]
        parts.append(f"File type: {file_type}")
        parts.append(f"Rows: {dataset.num_rows:,}")
        parts.append(f"Content format: {dataset.format}")

        if dataset.avg_tokens > 0:
            parts.append(
                f"Tokens: avg={dataset.avg_tokens:.0f}, "
                f"P50={dataset.p50_tokens:,}, P90={dataset.p90_tokens:,}, "
                f"P99={dataset.p99_tokens:,}, max={dataset.max_tokens:,}"
            )
            parts.append(f"Token method: {dataset.token_method}")

        if dataset.language:
            parts.append(f"Language: {dataset.language}")

        if dataset.duplicate_ratio > 0:
            parts.append(f"Duplicates: {dataset.duplicate_ratio:.1%}")

        if dataset.empty_ratio > 0:
            parts.append(f"Empty rows: {dataset.empty_ratio:.1%}")

        if dataset.text_column:
            parts.append(f"Text column: {dataset.text_column}")

        return " | ".join(parts)

    def _review_script_with_llm(
        self,
        script: str,
        config: TrainingConfig,
        hardware_summary: str,
        dataset_summary: str,
        llm_config: dict | None = None,
        file_type: str = "json",
        is_directory: bool = False,
    ) -> str:
        """Ask LLM to review and improve the generated script."""
        llm_config = llm_config or {}

        try:
            from bilgeai.llm import get_llm_client
            from bilgeai.llm.prompts import REVIEW_SCRIPT_PROMPT

            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 = REVIEW_SCRIPT_PROMPT.format(
                hardware_summary=hardware_summary,
                dataset_summary=dataset_summary,
                config_summary=json.dumps(config.model_dump(exclude={"explanations"}), indent=2),
                script=script,
                file_type=file_type,
            )

            response = client.generate(prompt)
            improved = response.content.strip()

            # Remove markdown fences if present
            if improved.startswith("```"):
                lines = improved.split("\n")
                improved = "\n".join(lines[1:-1]) if lines[-1].strip() == "```" else improved

            # Validate with ast.parse
            try:
                ast.parse(improved)
                # Post-fix: ensure load_dataset uses correct file type and param
                improved = _fix_load_dataset_file_type(improved, file_type, is_directory)
                logger.info("LLM improved the training script")
                return improved
            except SyntaxError:
                logger.warning("LLM returned invalid Python, keeping template output")

        except (LLMError, Exception):
            logger.debug("LLM script review failed, using template output", exc_info=True)

        return script
