"""Retrain runner: fine-tunes Llama 3.2 1B on merged training data.

Orchestrates the Unsloth + QLoRA fine-tuning pipeline:
  1. Validates training data
  2. Splits into train/eval (90/10, stratified by article)
  3. Writes and executes the Unsloth training script
  4. Evaluates against held-out set
  5. Exports to GGUF if eval passes

This module does NOT do the training itself (that requires GPU and
Unsloth in the Python environment). It generates the training script,
runs it via subprocess, and evaluates the result.

Usage:
    from air_blackbox.retrain.run import retrain, RetrainConfig
    result = retrain(RetrainConfig(training_data="training/training_data_v12.jsonl"))
    if result.success:
        print(f"Model at {result.model_path}, accuracy={result.eval_accuracy:.1%}")
"""

import json
import logging
import os
import random
import re
import subprocess
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

DEFAULT_OUTPUT_DIR = os.path.expanduser("~/.air-blackbox/models")
VALID_ARTICLES = {9, 10, 11, 12, 14, 15}


@dataclass
class RetrainConfig:
    """Configuration for a retrain run."""
    training_data: str                          # path to merged JSONL
    output_dir: str = DEFAULT_OUTPUT_DIR
    base_model: str = "unsloth/Llama-3.2-1B-bnb-4bit"
    epochs: int = 3
    batch_size: int = 4
    learning_rate: float = 2e-4
    lora_r: int = 16
    lora_alpha: int = 16
    max_seq_length: int = 2048
    eval_threshold: float = 0.90                # 90% accuracy minimum
    max_false_positive_rate: float = 0.05       # 5% max false positive rate
    max_false_negative_rate: float = 0.10       # 10% max false negative rate
    quantization: str = "q4_k_m"
    seed: int = 42
    eval_split: float = 0.10                    # 10% held out for eval


@dataclass
class RetrainResult:
    """Outcome of a retrain run."""
    success: bool = False
    model_path: str = ""
    gguf_path: str = ""
    training_loss: float = 0.0
    eval_accuracy: float = 0.0
    eval_false_positive_rate: float = 0.0
    eval_false_negative_rate: float = 0.0
    examples_trained: int = 0
    examples_eval: int = 0
    duration_seconds: int = 0
    version: str = ""
    error: str = ""
    eval_details: Dict = field(default_factory=dict)


def _extract_article(instruction: str) -> Optional[int]:
    m = re.search(r"Article\s+(\d+)", instruction)
    if m:
        num = int(m.group(1))
        return num if num in VALID_ARTICLES else None
    return None


def _stratified_split(
    examples: List[Dict], eval_fraction: float, seed: int
) -> tuple:
    """Split examples into train/eval, stratified by article."""
    rng = random.Random(seed)

    by_article: Dict[int, List[Dict]] = {}
    no_article: List[Dict] = []

    for ex in examples:
        article = _extract_article(ex.get("instruction", ""))
        if article is not None:
            by_article.setdefault(article, []).append(ex)
        else:
            no_article.append(ex)

    train, eval_set = [], []

    for article, items in by_article.items():
        rng.shuffle(items)
        n_eval = max(1, int(len(items) * eval_fraction))
        eval_set.extend(items[:n_eval])
        train.extend(items[n_eval:])

    # Non-article examples go to train.
    train.extend(no_article)

    rng.shuffle(train)
    rng.shuffle(eval_set)

    return train, eval_set


def _generate_training_script(
    config: RetrainConfig,
    train_path: str,
    output_path: str,
) -> str:
    """Generate the Unsloth fine-tuning script as a string."""
    return f'''#!/usr/bin/env python3
"""Auto-generated Unsloth fine-tuning script for AIR Blackbox compliance model.

Generated by air-blackbox retrain pipeline. Do not edit manually.
"""

import json
import torch
from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import Dataset

# Load training data.
examples = []
with open("{train_path}") as f:
    for line in f:
        line = line.strip()
        if line:
            examples.append(json.loads(line))

print(f"Loaded {{len(examples)}} training examples")

# Format for chat template.
def format_example(ex):
    return {{
        "text": (
            "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n"
            "You are an EU AI Act compliance analyzer. Analyze code for compliance gaps.\\n"
            "<|eot_id|><|start_header_id|>user<|end_header_id|>\\n"
            f"{{ex['instruction']}}\\n\\n{{ex['input']}}\\n"
            "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n"
            f"{{ex['output']}}<|eot_id|>"
        )
    }}

dataset = Dataset.from_list([format_example(ex) for ex in examples])

# Load model with QLoRA.
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="{config.base_model}",
    max_seq_length={config.max_seq_length},
    load_in_4bit=True,
    dtype=None,
)

model = FastLanguageModel.get_peft_model(
    model,
    r={config.lora_r},
    lora_alpha={config.lora_alpha},
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                     "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state={config.seed},
)

# Train.
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length={config.max_seq_length},
    args=TrainingArguments(
        output_dir="{output_path}",
        num_train_epochs={config.epochs},
        per_device_train_batch_size={config.batch_size},
        learning_rate={config.learning_rate},
        warmup_ratio=0.1,
        logging_steps=10,
        save_strategy="epoch",
        seed={config.seed},
        fp16=not torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False,
        bf16=torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False,
    ),
)

stats = trainer.train()

# Save training stats.
train_stats = {{
    "training_loss": stats.training_loss,
    "epochs": {config.epochs},
    "examples": len(examples),
    "global_steps": stats.global_step,
}}
with open("{output_path}/train_stats.json", "w") as f:
    json.dump(train_stats, f, indent=2)

print(f"Training complete. Loss: {{stats.training_loss:.4f}}")

# Merge LoRA weights and save.
model.save_pretrained_merged("{output_path}/merged", tokenizer, save_method="merged_16bit")
print("Merged model saved.")

# Export to GGUF.
model.save_pretrained_gguf(
    "{output_path}/gguf",
    tokenizer,
    quantization_method="{config.quantization}",
)
print("GGUF export complete.")
'''


def _run_eval(
    eval_examples: List[Dict],
    model_path: str,
    config: RetrainConfig,
) -> Dict:
    """Evaluate the fine-tuned model against held-out examples.

    Uses ollama to run inference on eval examples and compares
    the predicted article pass/fail status to the expected output.

    Returns dict with accuracy, false_positive_rate, false_negative_rate.
    """
    import shutil

    results = {
        "total": len(eval_examples),
        "correct": 0,
        "false_positives": 0,
        "false_negatives": 0,
        "errors": 0,
        "accuracy": 0.0,
        "false_positive_rate": 0.0,
        "false_negative_rate": 0.0,
    }

    # Check if ollama is available for eval.
    if not shutil.which("ollama"):
        logger.warning("Ollama not available, skipping model evaluation")
        results["error"] = "ollama not available"
        return results

    for ex in eval_examples:
        expected = ex.get("output", "")
        prompt = f"{ex['instruction']}\n\n{ex['input']}"

        try:
            proc = subprocess.run(
                ["ollama", "run", "air-compliance", prompt],
                capture_output=True,
                text=True,
                timeout=60,
            )
            predicted = proc.stdout.strip()

            # Compare pass/fail status per article.
            expected_passes = set(re.findall(r"Article\s+(\d+).*?PASS", expected))
            expected_fails = set(re.findall(r"Article\s+(\d+).*?FAIL", expected))
            predicted_passes = set(re.findall(r"Article\s+(\d+).*?PASS", predicted))
            predicted_fails = set(re.findall(r"Article\s+(\d+).*?FAIL", predicted))

            # A prediction is correct if pass/fail matches for the primary article.
            article = _extract_article(ex.get("instruction", ""))
            if article is not None:
                art_str = str(article)
                expected_status = "pass" if art_str in expected_passes else "fail"
                predicted_status = "pass" if art_str in predicted_passes else "fail"

                if expected_status == predicted_status:
                    results["correct"] += 1
                elif expected_status == "fail" and predicted_status == "pass":
                    results["false_negatives"] += 1  # missed a real issue
                elif expected_status == "pass" and predicted_status == "fail":
                    results["false_positives"] += 1  # cried wolf
            else:
                # Can't determine article, count as correct if outputs are similar.
                if any(word in predicted.lower() for word in ["pass", "fail", "finding"]):
                    results["correct"] += 1
                else:
                    results["errors"] += 1

        except (subprocess.TimeoutExpired, Exception) as e:
            results["errors"] += 1
            logger.debug(f"Eval error: {e}")

    assessed = results["total"] - results["errors"]
    if assessed > 0:
        results["accuracy"] = results["correct"] / assessed
        results["false_positive_rate"] = results["false_positives"] / assessed
        results["false_negative_rate"] = results["false_negatives"] / assessed

    return results


def retrain(config: RetrainConfig) -> RetrainResult:
    """Run the full retrain pipeline.

    1. Load and split training data
    2. Generate and execute Unsloth training script
    3. Evaluate against held-out set
    4. Export to GGUF if eval passes

    Args:
        config: RetrainConfig with paths and hyperparameters

    Returns:
        RetrainResult with success status, model path, and eval metrics
    """
    result = RetrainResult()
    start_time = time.time()

    training_data = os.path.expanduser(config.training_data)
    output_dir = os.path.expanduser(config.output_dir)

    # Detect version from filename.
    m = re.search(r"v(\d+)", os.path.basename(training_data))
    result.version = f"v{m.group(1)}" if m else "v0"

    # Create version-specific output directory.
    version_dir = os.path.join(output_dir, result.version)
    os.makedirs(version_dir, exist_ok=True)

    # 1. Load training data.
    examples = []
    with open(training_data, "r") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    examples.append(json.loads(line))
                except json.JSONDecodeError:
                    continue

    if not examples:
        result.error = "No valid training examples found"
        return result

    # 2. Stratified split.
    train, eval_set = _stratified_split(examples, config.eval_split, config.seed)
    result.examples_trained = len(train)
    result.examples_eval = len(eval_set)

    logger.info(
        f"Split: {len(train)} train, {len(eval_set)} eval "
        f"from {len(examples)} total"
    )

    # Write split files.
    train_path = os.path.join(version_dir, "train.jsonl")
    eval_path = os.path.join(version_dir, "eval.jsonl")

    with open(train_path, "w") as f:
        for ex in train:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")
    with open(eval_path, "w") as f:
        for ex in eval_set:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    # 3. Generate training script.
    script_content = _generate_training_script(config, train_path, version_dir)
    script_path = os.path.join(version_dir, "train.py")
    with open(script_path, "w") as f:
        f.write(script_content)

    logger.info(f"Training script written to {script_path}")

    # 4. Execute training.
    try:
        proc = subprocess.run(
            ["python3", script_path],
            capture_output=True,
            text=True,
            timeout=7200,  # 2 hour timeout
            cwd=version_dir,
        )

        if proc.returncode != 0:
            result.error = f"Training failed: {proc.stderr[-500:]}"
            result.duration_seconds = int(time.time() - start_time)
            # Save error log.
            with open(os.path.join(version_dir, "error.log"), "w") as f:
                f.write(proc.stderr)
            return result

    except subprocess.TimeoutExpired:
        result.error = "Training timed out after 2 hours"
        result.duration_seconds = int(time.time() - start_time)
        return result

    # 5. Read training stats.
    stats_path = os.path.join(version_dir, "train_stats.json")
    if os.path.exists(stats_path):
        with open(stats_path) as f:
            stats = json.load(f)
        result.training_loss = stats.get("training_loss", 0)

    # 6. Check GGUF output.
    gguf_dir = os.path.join(version_dir, "gguf")
    gguf_files = []
    if os.path.isdir(gguf_dir):
        gguf_files = [f for f in os.listdir(gguf_dir) if f.endswith(".gguf")]

    if not gguf_files:
        result.error = "GGUF export failed: no .gguf file found"
        result.duration_seconds = int(time.time() - start_time)
        return result

    result.gguf_path = os.path.join(gguf_dir, gguf_files[0])
    result.model_path = version_dir

    # 7. Evaluate (if ollama is available).
    eval_results = _run_eval(eval_set, version_dir, config)
    result.eval_accuracy = eval_results["accuracy"]
    result.eval_false_positive_rate = eval_results["false_positive_rate"]
    result.eval_false_negative_rate = eval_results["false_negative_rate"]
    result.eval_details = eval_results

    # 8. Check thresholds.
    if eval_results.get("error"):
        # Eval couldn't run (no ollama), pass with warning.
        logger.warning(
            f"Evaluation skipped ({eval_results['error']}). "
            "Model exported but not validated."
        )
        result.success = True
    elif result.eval_accuracy < config.eval_threshold:
        result.error = (
            f"Eval accuracy {result.eval_accuracy:.1%} below threshold "
            f"{config.eval_threshold:.1%}. Model saved but NOT published."
        )
        result.success = False
    elif result.eval_false_positive_rate > config.max_false_positive_rate:
        result.error = (
            f"False positive rate {result.eval_false_positive_rate:.1%} above "
            f"threshold {config.max_false_positive_rate:.1%}. Model NOT published."
        )
        result.success = False
    else:
        result.success = True

    result.duration_seconds = int(time.time() - start_time)

    logger.info(
        "retrain_complete",
        extra={
            "success": result.success,
            "accuracy": result.eval_accuracy,
            "fp_rate": result.eval_false_positive_rate,
            "loss": result.training_loss,
            "duration": result.duration_seconds,
        },
    )

    return result
