"""LLM-specific pipeline: sensitivity analysis, optimization, conversion."""

import os
import json
from typing import Optional

from ..core.importer import load_hf_model, ModelInfo
from ..core.sensitivity import (
    analyze_llm_sensitivity,
    print_sensitivity_report,
    SensitivityResult,
)
from ..core.optimizer import (
    optimize_mixed_precision,
    make_quant_predicate,
    OptimizationResult,
)
from ..backends.mlx_backend import (
    convert_llm_to_mlx,
    convert_llm_uniform,
    convert_llm_static_mixed,
)
from ..calibration.datasets import load_llm_calibration


def run_llm_pipeline(
    model_name: str,
    output_dir: str,
    target_bpw: float = 4.5,
    n_calibration: int = 8,
    candidate_bits: list[int] | None = None,
    group_size: int = 64,
    skip_baselines: bool = False,
    # Legacy compat
    low_bits: int = 4,
    high_bits: int = 8,
) -> dict:
    """Run the full OptiQ LLM pipeline.

    1. Load model in PyTorch
    2. Run per-layer KL divergence sensitivity analysis
    3. Multi-tier optimization (greedy knapsack)
    4. Convert to MLX (OptiQ + baselines)

    Returns dict with output paths and results.
    """
    if candidate_bits is None:
        candidate_bits = [4, 8]

    os.makedirs(output_dir, exist_ok=True)
    results = {}

    # Step 1: Load model
    print("\n[1/4] Loading model...")
    model_info = load_hf_model(model_name)

    # Step 2: Per-layer KL divergence sensitivity analysis
    print(f"\n[2/4] Running per-layer KL sensitivity (bits: {candidate_bits})...")
    calibration_data = load_llm_calibration(
        model_info.tokenizer,
        n_samples=max(n_calibration * 4, 32),
        seq_len=128,
    )
    checkpoint_path = os.path.join(output_dir, "sensitivity_checkpoint.json")
    sensitivity_results = analyze_llm_sensitivity(
        model_info.model,
        calibration_data,
        candidate_bits=candidate_bits,
        group_size=group_size,
        n_calibration=min(n_calibration, 2),
        seq_len=128,
        checkpoint_path=checkpoint_path,
    )
    print_sensitivity_report(sensitivity_results)
    results["sensitivity"] = sensitivity_results

    # Step 3: Multi-tier optimization
    print("\n[3/4] Optimizing multi-tier bit allocation...")
    opt_result = optimize_mixed_precision(
        sensitivity_results,
        target_bpw=target_bpw,
        candidate_bits=candidate_bits,
        group_size=group_size,
    )
    results["optimization"] = opt_result

    # Step 4: Convert
    print("\n[4/4] Converting to MLX...")

    # OptiQ mixed-precision
    optiq_path = os.path.join(output_dir, "optiq_mixed")
    predicate = make_quant_predicate(opt_result)
    metadata = {
        "method": "optiq_mixed_precision",
        "target_bpw": target_bpw,
        "achieved_bpw": opt_result.achieved_bpw,
        "n_high_bits": opt_result.n_high_bits,
        "n_low_bits": opt_result.n_low_bits,
        "threshold": opt_result.threshold,
        "per_layer": {
            c.layer_name: {"bits": c.bits, "group_size": c.group_size}
            for c in opt_result.configs
        },
    }
    convert_llm_to_mlx(
        model_name,
        optiq_path,
        quant_predicate=predicate,
        q_bits=low_bits,
        q_group_size=group_size,
        metadata=metadata,
    )
    results["optiq_path"] = optiq_path

    # Baselines
    if not skip_baselines:
        # Uniform 4-bit
        uniform_path = os.path.join(output_dir, "uniform_4bit")
        convert_llm_uniform(model_name, uniform_path, bits=4, group_size=group_size)
        results["uniform_path"] = uniform_path

        # Static mixed (mlx-lm recipe)
        try:
            static_path = os.path.join(output_dir, "static_mixed")
            convert_llm_static_mixed(model_name, static_path, recipe="mixed_3_6")
            results["static_path"] = static_path
        except Exception as e:
            print(f"  Static mixed recipe failed (may not support this model): {e}")

    # Free PyTorch model
    del model_info
    import gc
    gc.collect()

    return results


def load_sensitivity_cache(cache_path: str) -> Optional[list[SensitivityResult]]:
    """Load cached sensitivity results."""
    if not os.path.exists(cache_path):
        return None
    with open(cache_path) as f:
        data = json.load(f)
    return [
        SensitivityResult(
            layer_name=d["layer_name"],
            sensitivities={int(k): v for k, v in d["sensitivities"].items()},
            param_count=d["param_count"],
        )
        for d in data
    ]


def save_sensitivity_cache(results: list[SensitivityResult], cache_path: str):
    """Save sensitivity results to cache."""
    data = [
        {
            "layer_name": r.layer_name,
            "sensitivities": {str(k): v for k, v in r.sensitivities.items()},
            "param_count": r.param_count,
        }
        for r in results
    ]
    with open(cache_path, "w") as f:
        json.dump(data, f, indent=2)
