"""Code generator: renders training scripts from TrainingConfig."""

from __future__ import annotations

import ast
import json
import logging
from pathlib import Path

from jinja2 import Environment, FileSystemLoader

from bilgeai.optimizer.engine import _detect_file_type
from bilgeai.optimizer.models import TrainingConfig

TEMPLATES_DIR = Path(__file__).parent / "templates"

logger = logging.getLogger(__name__)

# Map Python top-level module names to pip package names
_MODULE_TO_PACKAGE: dict[str, str] = {
    "torch": "torch",
    "transformers": "transformers",
    "datasets": "datasets",
    "trl": "trl",
    "peft": "peft",
    "accelerate": "accelerate",
    "bitsandbytes": "bitsandbytes",
    "deepspeed": "deepspeed",
    "numpy": "numpy",
    "wandb": "wandb",
    "mlflow": "mlflow",
    "tensorboard": "tensorboard",
    "scipy": "scipy",
    "sklearn": "scikit-learn",
    "sentencepiece": "sentencepiece",
    "safetensors": "safetensors",
}

# Standard library modules to ignore
_STDLIB = {
    "random",
    "os",
    "sys",
    "json",
    "math",
    "pathlib",
    "logging",
    "typing",
    "collections",
    "functools",
    "itertools",
    "abc",
    "re",
    "copy",
    "time",
    "datetime",
    "dataclasses",
    "enum",
    "io",
    "csv",
    "argparse",
    "subprocess",
    "shutil",
    "glob",
    "hashlib",
    "secrets",
    "warnings",
    "contextlib",
    "textwrap",
}


def extract_dependencies(script: str) -> list[str]:
    """Parse generated Python script and return required pip packages."""
    try:
        tree = ast.parse(script)
    except SyntaxError:
        logger.debug("Failed to parse script for dependency extraction")
        return []

    modules: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                modules.add(alias.name.split(".")[0])
        elif isinstance(node, ast.ImportFrom) and node.module:
            modules.add(node.module.split(".")[0])

    # Filter out stdlib and map to pip names
    packages: list[str] = []
    seen: set[str] = set()
    for mod in sorted(modules):
        if mod in _STDLIB:
            continue
        pkg = _MODULE_TO_PACKAGE.get(mod, mod)
        if pkg not in seen:
            seen.add(pkg)
            packages.append(pkg)

    return packages


def _format_code(code: str) -> str:
    """Format generated Python code with black if available, else return as-is."""
    try:
        import black

        mode = black.Mode(line_length=100)
        return black.format_str(code, mode=mode)
    except ImportError:
        logger.debug("black not installed, skipping code formatting")
        return code
    except Exception:
        logger.debug("black formatting failed, returning unformatted code", exc_info=True)
        return code


class CodeGenerator:
    """Generates training scripts from TrainingConfig using Jinja2 templates."""

    def __init__(self) -> None:
        self.env = Environment(
            loader=FileSystemLoader(str(TEMPLATES_DIR)),
            keep_trailing_newline=True,
            trim_blocks=True,
            lstrip_blocks=True,
        )

    def generate_sft(
        self,
        config: TrainingConfig,
        dataset_path: str = "./train.jsonl",
        dataset_format: str = "alpaca",
        hardware_summary: str = "",
    ) -> str:
        """Generate an SFT+LoRA training script from config."""
        template = self.env.get_template("sft_lora.py.j2")

        # Detect file type from extension
        file_type = _detect_file_type(dataset_path)
        is_directory = Path(dataset_path).is_dir() or not Path(dataset_path).suffix

        context = {
            **config.model_dump(exclude={"explanations"}),
            "dataset_path": dataset_path,
            "dataset_format": dataset_format,
            "file_type": file_type,
            "is_directory": is_directory,
            "hardware_summary": hardware_summary,
            "deepspeed_config_dict": self._build_deepspeed_config(config),
        }

        return _format_code(template.render(**context))

    def generate_dpo(
        self,
        config: TrainingConfig,
        dataset_path: str = "./train.jsonl",
        hardware_summary: str = "",
    ) -> str:
        """Generate a DPO training script from config."""
        template = self.env.get_template("dpo.py.j2")

        context = {
            **config.model_dump(exclude={"explanations"}),
            "dataset_path": dataset_path,
            "hardware_summary": hardware_summary,
        }

        return _format_code(template.render(**context))

    def generate_rlhf(
        self,
        config: TrainingConfig,
        dataset_path: str = "./train.jsonl",
        dataset_format: str = "alpaca",
        hardware_summary: str = "",
        reward_model_name: str | None = None,
        ppo_epochs: int = 4,
    ) -> str:
        """Generate an RLHF (PPO) training script from config."""
        template = self.env.get_template("rlhf.py.j2")

        context = {
            **config.model_dump(exclude={"explanations"}),
            "dataset_path": dataset_path,
            "dataset_format": dataset_format,
            "hardware_summary": hardware_summary,
            "reward_model_name": reward_model_name or config.model_name,
            "ppo_epochs": ppo_epochs,
        }

        return _format_code(template.render(**context))

    def generate_and_save(
        self,
        config: TrainingConfig,
        output_path: Path,
        dataset_path: str = "./train.jsonl",
        dataset_format: str = "alpaca",
        hardware_summary: str = "",
    ) -> Path:
        """Generate script and write to file."""
        code = self.generate_sft(
            config,
            dataset_path=dataset_path,
            dataset_format=dataset_format,
            hardware_summary=hardware_summary,
        )
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(code, encoding="utf-8")
        return output_path

    @staticmethod
    def _build_deepspeed_config(config: TrainingConfig) -> str:
        """Build a DeepSpeed config dict string for template injection."""
        if config.distributed_strategy != "deepspeed" or config.deepspeed_stage is None:
            return "{}"

        ds_config: dict = {
            "zero_optimization": {
                "stage": config.deepspeed_stage,
                "overlap_comm": True,
                "contiguous_gradients": True,
            },
            "gradient_accumulation_steps": config.gradient_accumulation_steps,
            "train_micro_batch_size_per_gpu": config.batch_size,
        }

        if config.deepspeed_offload_optimizer:
            ds_config["zero_optimization"]["offload_optimizer"] = {
                "device": "cpu",
                "pin_memory": True,
            }

        if config.deepspeed_offload_param:
            ds_config["zero_optimization"]["offload_param"] = {
                "device": "cpu",
                "pin_memory": True,
            }

        if config.bf16:
            ds_config["bf16"] = {"enabled": True}
        elif config.fp16:
            ds_config["fp16"] = {"enabled": True}

        return json.dumps(ds_config, indent=4)
