"""
Modal app: Style-SFT training for BART-large on A100 80GB.

Trains BARTWithStyleEmbeddings on HC3 AI→Human parallel corpus.
Output: fine-tuned model checkpoint uploaded to HuggingFace.

Budget: ~15-25€, max 8h A100 80GB.

Usage:
  # Dry-run (50 steps, local validation, 0€)
  modal run src/modal_app_sft.py --dry-run

  # Real training on A100
  modal run src/modal_app_sft.py --data datasets/style_transfer_pairs_train.jsonl
"""

from __future__ import annotations

import json
import os
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import modal

# ---------------------------------------------------------------------------
# Modal image with training dependencies
# ---------------------------------------------------------------------------

image = (
    modal.Image.debian_slim(python_version="3.12")
    .env({"PIP_PROGRESS_BAR": "off", "PYTHONIOENCODING": "utf-8"})
    .pip_install(
        "torch>=2.4.0",
        "transformers>=4.45.0",
        "accelerate>=0.34.0",
        "datasets>=3.0.0",
        "numpy>=1.26.0",
    )
)

app = modal.App("evasion-detection-sft", image=image)

# HF cache volume for model weights
hf_cache = modal.Volume.from_name("hf-cache", create_if_missing=True)


# ---------------------------------------------------------------------------
# CostGuard (from shared module — no torch dependency)
# ---------------------------------------------------------------------------

@dataclass
class CostGuard:
    max_runtime_hours: float = 8.0
    max_cost_eur: float = 25.0
    gpu_type: str = "A100-80GB"
    dry_run: bool = True
    dry_run_max_steps: int = 50

    def validate(self, elapsed_hours: float) -> bool:
        rates = {"T4": 0.60, "L4": 0.80, "A10G": 1.10, "A100-80GB": 2.50, "H100": 3.95}
        rate = rates.get(self.gpu_type, 2.50)
        if elapsed_hours > self.max_runtime_hours:
            print(f"[CostGuard] TIMEOUT: {elapsed_hours:.1f}h > {self.max_runtime_hours}h")
            return False
        cost = elapsed_hours * rate
        if cost > self.max_cost_eur:
            print(f"[CostGuard] OVER BUDGET: {cost:.2f}EUR > {self.max_cost_eur}EUR")
            return False
        return True


# ---------------------------------------------------------------------------
# Training config
# ---------------------------------------------------------------------------

@dataclass
class TrainingConfig:
    model_name: str = "facebook/bart-large"
    output_dir: str = "/tmp/bart_sft_output"
    batch_size: int = 8
    gradient_accumulation_steps: int = 4     # effective batch = 32
    learning_rate: float = 2e-5
    warmup_steps: int = 200
    max_steps: int = 5000
    eval_steps: int = 500
    save_steps: int = 1000
    max_input_length: int = 512
    max_output_length: int = 512
    seed: int = 42
    fp16: bool = True
    style_id_ai: int = 0
    style_id_human: int = 1
    # HF upload
    hf_repo: str = "simonlesaumon/evasion-detection-models"
    hf_model_name: str = "bart-sft-style-humanization"


# ---------------------------------------------------------------------------
# Dataset (runs inside Modal container)
# ---------------------------------------------------------------------------

def load_training_data(data_path: str, tokenizer, max_input_length: int, max_output_length: int,
                       style_token_ai: str = "<ai>", style_token_human: str = "<human>"):
    """Load style transfer pairs from HF dataset path or local JSONL."""
    import torch
    from torch.utils.data import Dataset

    class StyleTransferDataset(Dataset):
        def __init__(self, samples, tokenizer, max_input_length, max_output_length,
                     style_ai, style_human):
            self.tokenizer = tokenizer
            self.max_input_length = max_input_length
            self.max_output_length = max_output_length
            self.style_ai = style_ai
            self.style_human = style_human
            self.samples = samples

            # Add style tokens to vocabulary
            for tok in [style_ai, style_human]:
                if tok not in tokenizer.get_vocab():
                    tokenizer.add_tokens([tok])

        def __len__(self):
            return len(self.samples)

        def __getitem__(self, idx):
            sample = self.samples[idx]
            # Input: <human> AI_text → Output: human_text
            input_text = f"{self.style_human} {sample['ai_text']}"
            target_text = sample["human_text"]

            inputs = self.tokenizer(
                input_text,
                max_length=self.max_input_length,
                truncation=True,
                padding="max_length",
                return_tensors="pt",
            )

            targets = self.tokenizer(
                target_text,
                max_length=self.max_output_length,
                truncation=True,
                padding="max_length",
                return_tensors="pt",
            )

            return {
                "input_ids": inputs["input_ids"].squeeze(0),
                "attention_mask": inputs["attention_mask"].squeeze(0),
                "labels": targets["input_ids"].squeeze(0),
            }

    # Load samples
    samples = []
    if data_path.startswith("datasets/"):
        # Download raw JSONL from HF repo directly (not via datasets library)
        import requests
        hf_token = os.getenv("HF_TOKEN", "")
        headers = {}
        if hf_token:
            headers["Authorization"] = f"Bearer {hf_token}"

        url = f"https://huggingface.co/simonlesaumon/evasion-detection-artifacts/resolve/main/{data_path}"
        print(f"[SFT] Downloading data from {url}...")
        resp = requests.get(url, headers=headers, timeout=60)

        if resp.status_code == 200:
            for line in resp.text.splitlines():
                if line.strip():
                    try:
                        samples.append(json.loads(line))
                    except json.JSONDecodeError:
                        continue
            print(f"[SFT] Downloaded {len(samples)} samples from HF")
        else:
            print(f"[SFT] HTTP {resp.status_code} downloading data, using fallback")
            samples = _get_fallback_samples()
    elif os.path.exists(data_path):
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    samples.append(json.loads(line))
        print(f"[SFT] Loaded {len(samples)} samples from {data_path}")
    else:
        # Fallback: use synthetic data
        print(f"[SFT] WARNING: Data not found at {data_path}, using synthetic fallback")
        samples = _get_fallback_samples()

    print(f"[SFT] Loaded {len(samples)} training samples")
    return StyleTransferDataset(samples, tokenizer, max_input_length, max_output_length,
                                style_token_ai, style_token_human)


def _get_fallback_samples() -> list[dict]:
    """Synthetic fallback pairs if dataset not available."""
    return [
        {"ai_text": "The implementation of machine learning algorithms has demonstrated "
                     "significant improvements in various domains. These systems leverage "
                     "large datasets to identify patterns and make predictions with high accuracy.",
         "human_text": "So I tried using ML for this project and honestly it worked way better "
                       "than I expected. You feed it a bunch of data and it somehow figures out "
                       "patterns you'd never spot manually.",
         "domain": "tech"},
    ] * 100  # Repeat to have enough for training


# ---------------------------------------------------------------------------
# Model with Style Embeddings
# ---------------------------------------------------------------------------

def create_style_model(model_name: str = "facebook/bart-large"):
    """Create BART model with trainable style embeddings."""
    import torch
    import torch.nn as nn
    from transformers import AutoModelForSeq2SeqLM

    class BARTWithStyleEmbeddings(nn.Module):
        def __init__(self, model_name="facebook/bart-large", style_dim=1024, num_styles=2):
            super().__init__()
            self.bart = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            self.config = self.bart.config
            hidden_size = self.config.d_model
            self.style_embeddings = nn.Embedding(num_styles, hidden_size)
            self.style_proj = nn.Sequential(
                nn.Linear(hidden_size, style_dim),
                nn.GELU(),
                nn.Linear(style_dim, hidden_size),
            )
            nn.init.normal_(self.style_embeddings.weight, std=0.02)
            self._style_injection = None  # cache for generation

        def set_style(self, style_id: int):
            """Set the style to inject during generation."""
            style_emb = self.style_embeddings(torch.tensor([style_id]))
            self._style_injection = self.style_proj(style_emb)  # (1, hidden)

        def forward(self, input_ids, attention_mask=None, labels=None, style_ids=None):
            encoder_outputs = self.bart.model.encoder(
                input_ids=input_ids, attention_mask=attention_mask, return_dict=True,
            )
            if style_ids is not None:
                style_emb = self.style_embeddings(style_ids)
                style_emb = self.style_proj(style_emb)
                encoder_outputs.last_hidden_state = (
                    encoder_outputs.last_hidden_state + style_emb.unsqueeze(1)
                )
            decoder_outputs = self.bart(
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                labels=labels,
                return_dict=True,
            )
            return decoder_outputs

        def generate(self, input_ids, attention_mask=None, max_length=512, **kwargs):
            """Generate with cached style injection via encoder outputs modification."""
            encoder_outputs = self.bart.model.encoder(
                input_ids=input_ids, attention_mask=attention_mask, return_dict=True,
            )
            if self._style_injection is not None:
                style_emb = self._style_injection.to(encoder_outputs.last_hidden_state.device)
                encoder_outputs.last_hidden_state = (
                    encoder_outputs.last_hidden_state + style_emb.unsqueeze(1)
                )
            return self.bart.generate(
                encoder_outputs=encoder_outputs,
                attention_mask=attention_mask,
                max_length=max_length,
                **kwargs,
            )

        def save_pretrained(self, path: str):
            os.makedirs(path, exist_ok=True)
            self.bart.save_pretrained(path)
            torch.save({
                "style_embeddings": self.style_embeddings.state_dict(),
                "style_proj": self.style_proj.state_dict(),
            }, os.path.join(path, "style_modules.pt"))

        @classmethod
        def from_pretrained(cls, path: str, model_name: str = "facebook/bart-large"):
            import torch as t
            model = cls(model_name)
            model.bart = AutoModelForSeq2SeqLM.from_pretrained(path)
            ckpt = t.load(os.path.join(path, "style_modules.pt"), map_location="cpu")
            model.style_embeddings.load_state_dict(ckpt["style_embeddings"])
            model.style_proj.load_state_dict(ckpt["style_proj"])
            return model

    return BARTWithStyleEmbeddings(model_name)


# ---------------------------------------------------------------------------
# Training loop (runs inside Modal container)
# ---------------------------------------------------------------------------

def train_sft_on_gpu(
    data_path: str,
    config: TrainingConfig,
    cost_guard: CostGuard,
) -> dict:
    """Run Style-SFT training on GPU."""
    import torch
    from transformers import AutoTokenizer, get_linear_schedule_with_warmup
    from torch.utils.data import DataLoader

    torch.manual_seed(config.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[SFT] Device: {device}")
    print(f"[SFT] Config: {config}")

    # Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(config.model_name)

    # Dataset
    dataset = load_training_data(
        data_path, tokenizer, config.max_input_length, config.max_output_length
    )

    if cost_guard.dry_run:
        print(f"[SFT] DRY RUN: limiting to {cost_guard.dry_run_max_steps} steps")
        dataset.samples = dataset.samples[: min(20, len(dataset.samples))]
        config.max_steps = cost_guard.dry_run_max_steps

    dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True, num_workers=2)

    # Model
    print(f"[SFT] Creating model: {config.model_name}")
    BARTModel = create_style_model(config.model_name)
    model = BARTModel
    model.to(device)

    # Resize token embeddings if style tokens added
    if len(tokenizer) > model.config.vocab_size:
        model.bart.resize_token_embeddings(len(tokenizer))

    # Optimizer + scheduler
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=config.warmup_steps,
        num_training_steps=config.max_steps,
    )

    # Training
    model.train()
    global_step = 0
    total_loss = 0.0
    best_loss = float("inf")
    start_time = time.time()

    print(f"[SFT] Starting training: {config.max_steps} steps, "
          f"batch={config.batch_size}, grad_accum={config.gradient_accumulation_steps}")

    scaler = torch.amp.GradScaler("cuda") if config.fp16 and device.type == "cuda" else None

    for epoch in range(10):
        for batch in dataloader:
            # Cost guard check every 100 steps
            if global_step % 100 == 0:
                elapsed_h = (time.time() - start_time) / 3600
                if not cost_guard.validate(elapsed_h):
                    return {"status": "aborted_cost_guard", "step": global_step}

            batch = {k: v.to(device) for k, v in batch.items()}
            style_ids = torch.full(
                (batch["input_ids"].shape[0],), config.style_id_human,
                dtype=torch.long, device=device,
            )

            if scaler is not None:
                with torch.amp.autocast("cuda"):
                    outputs = model(
                        input_ids=batch["input_ids"],
                        attention_mask=batch["attention_mask"],
                        labels=batch["labels"],
                        style_ids=style_ids,
                    )
                    loss = outputs.loss / config.gradient_accumulation_steps
                scaler.scale(loss).backward()
            else:
                outputs = model(
                    input_ids=batch["input_ids"],
                    attention_mask=batch["attention_mask"],
                    labels=batch["labels"],
                    style_ids=style_ids,
                )
                loss = outputs.loss / config.gradient_accumulation_steps
                loss.backward()

            total_loss += loss.item()

            if (global_step + 1) % config.gradient_accumulation_steps == 0:
                if scaler is not None:
                    scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                if scaler is not None:
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

            global_step += 1

            # Logging
            if global_step % config.eval_steps == 0:
                avg_loss = total_loss / config.eval_steps
                elapsed_h = (time.time() - start_time) / 3600
                cost = elapsed_h * 2.50
                print(f"[SFT] Step {global_step}/{config.max_steps} | "
                      f"Loss: {avg_loss:.4f} | "
                      f"Time: {elapsed_h:.1f}h | "
                      f"Cost: ~{cost:.2f}EUR | "
                      f"LR: {scheduler.get_last_lr()[0]:.2e}")
                total_loss = 0.0

                if avg_loss < best_loss:
                    best_loss = avg_loss

            # Checkpoint
            if global_step % config.save_steps == 0:
                ckpt_path = os.path.join(config.output_dir, f"checkpoint-{global_step}")
                model.save_pretrained(ckpt_path)
                tokenizer.save_pretrained(ckpt_path)
                print(f"[SFT] Checkpoint saved: {ckpt_path}")

            if global_step >= config.max_steps:
                break

        if global_step >= config.max_steps:
            break

    # Final save
    final_path = os.path.join(config.output_dir, "final")
    model.save_pretrained(final_path)
    tokenizer.save_pretrained(final_path)

    elapsed_h = (time.time() - start_time) / 3600
    result = {
        "status": "completed",
        "total_steps": global_step,
        "best_loss": best_loss,
        "elapsed_hours": round(elapsed_h, 2),
        "estimated_cost_eur": round(elapsed_h * 2.50, 2),
        "final_model_path": final_path,
        "config": asdict(config),
    }
    print(f"[SFT] Training done: {json.dumps(result, indent=2)}")
    return result


# ---------------------------------------------------------------------------
# HuggingFace upload (inside Modal container)
# ---------------------------------------------------------------------------

def upload_to_hf(local_dir: str, repo_id: str, model_name: str) -> str:
    """Upload trained model to HuggingFace."""
    import subprocess

    hf_token = os.getenv("HF_TOKEN", "")
    if not hf_token:
        print("[SFT] WARNING: HF_TOKEN not set. Skipping upload.")
        return "skipped: no HF_TOKEN"

    print(f"[SFT] Uploading model to {repo_id}/{model_name}...")
    try:
        # Upload all files in the final checkpoint
        subprocess.run([
            "huggingface-cli", "upload", repo_id, local_dir, model_name,
            "--token", hf_token,
        ], check=True)
        print(f"[SFT] Uploaded to https://huggingface.co/{repo_id}")
        return f"https://huggingface.co/{repo_id}/tree/main/{model_name}"
    except Exception as e:
        print(f"[SFT] Upload failed: {e}")
        return f"upload_failed: {e}"


# ---------------------------------------------------------------------------
# Modal function
# ---------------------------------------------------------------------------

@app.function(
    gpu="A100-80GB",
    timeout=60 * 60 * 10,  # 10h max
    scaledown_window=60 * 10,
    volumes={"/root/.cache/huggingface": hf_cache},
)
def train_style_sft(
    data_path: str = "datasets/style_transfer_pairs_train.jsonl",
    model_name: str = "facebook/bart-large",
    max_steps: int = 5000,
    batch_size: int = 8,
    learning_rate: float = 2e-5,
    dry_run: bool = False,
    push_to_hf: bool = True,
) -> dict:
    """Run Style-SFT training on A100 80GB."""
    config = TrainingConfig(
        model_name=model_name,
        max_steps=max_steps,
        batch_size=batch_size,
        learning_rate=learning_rate,
    )

    guard = CostGuard(
        dry_run=dry_run,
        gpu_type="A100-80GB",
    )

    result = train_sft_on_gpu(data_path, config, guard)

    if result["status"] == "completed" and push_to_hf and not dry_run:
        hf_url = upload_to_hf(
            os.path.join(config.output_dir, "final"),
            config.hf_repo,
            config.hf_model_name,
        )
        result["hf_url"] = hf_url

    # Save result metadata
    os.makedirs(config.output_dir, exist_ok=True)
    with open(os.path.join(config.output_dir, "train_result.json"), "w") as f:
        json.dump(result, f, indent=2)

    return result


# ---------------------------------------------------------------------------
# Local entrypoint
# ---------------------------------------------------------------------------

@app.local_entrypoint()
def main(
    data: str = "datasets/style_transfer_pairs_train.jsonl",
    model: str = "facebook/bart-large",
    max_steps: int = 5000,
    batch_size: int = 8,
    lr: float = 2e-5,
    dry_run: bool = False,
    push_to_hf: bool = False,
):
    """Local entrypoint — dispatches training to Modal A100."""
    print("=" * 60)
    print(" Style-SFT Training — Evasion Detection")
    print("=" * 60)
    print(f"  Model: {model}")
    print(f"  Data: {data}")
    print(f"  Steps: {max_steps}")
    print(f"  Batch: {batch_size}")
    print(f"  LR: {lr}")
    print(f"  Dry-run: {dry_run}")
    print(f"  Push to HF: {push_to_hf}")
    print("=" * 60)

    if dry_run:
        print("\n[SFT] DRY RUN — validating pipeline (no A100, ~50 steps on T4)...")
        result = train_style_sft.local(
            data_path=data,
            model_name=model,
            max_steps=50,
            batch_size=2,
            learning_rate=lr,
            dry_run=True,
            push_to_hf=False,
        )
    else:
        print("\n[SFT] Launching training on Modal A100-80GB (~15-25€, 6-8h)...")
        result = train_style_sft.remote(
            data_path=data,
            model_name=model,
            max_steps=max_steps,
            batch_size=batch_size,
            learning_rate=lr,
            dry_run=False,
            push_to_hf=push_to_hf,
        )

    print(f"\n[SFT] Result: {json.dumps(result, indent=2, default=str)}")

    if result.get("status") == "completed":
        print(f"\n[SFT] Training complete! Best loss: {result.get('best_loss', 'N/A')}")
        print(f"[SFT] Time: {result.get('elapsed_hours', 'N/A')}h")
        print(f"[SFT] Cost: ~{result.get('estimated_cost_eur', 'N/A')}EUR")
        if result.get("hf_url"):
            print(f"[SFT] Model: {result['hf_url']}")

    return result