"""
Modal app: CoPA (Contrastive Paraphrase Attack) on GPU.

Deploys the CoPA prototype to Modal cloud — no local torch needed.
Dry-run mode: just validates the pipeline without launching real GPU.

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

  # Real run on T4 GPU (~0.60€/h)
  modal run src/modal_app_copa.py

  # Real run on A100 80GB (~2.50€/h)
  modal run src/modal_app_copa.py --gpu A100-80GB
"""

from __future__ import annotations

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

import modal

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

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

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


# ---------------------------------------------------------------------------
# CostGuard (lightweight, no deps beyond stdlib)
# ---------------------------------------------------------------------------

@dataclass
class CostGuard:
    max_runtime_minutes: int = 30
    max_cost_eur: float = 5.0
    gpu_type: str = "T4"
    dry_run: bool = True

    def validate(self, elapsed_m: 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_m > self.max_runtime_minutes:
            print(f"[CostGuard] TIMEOUT: {elapsed_m:.0f}m > {self.max_runtime_minutes}m")
            return False
        cost = (elapsed_m / 60) * rate
        if cost > self.max_cost_eur:
            print(f"[CostGuard] OVER BUDGET: {cost:.2f}€ > {self.max_cost_eur}€")
            return False
        return True


# ---------------------------------------------------------------------------
# CoPA logic (runs INSIDE Modal container, so torch is available)
# ---------------------------------------------------------------------------

def _run_copa_on_gpu(
    texts: list[str],
    model_name: str = "Qwen/Qwen2.5-1.5B-Instruct",
    lambda_contrast: float = 0.8,
    alpha_truncation: float = 1e-5,
    max_new_tokens: int = 768,
    temperature: float = 1.0,
    top_p: float = 0.92,
    repetition_penalty: float = 1.15,
    diversity_bonus_strength: float = 0.5,
    human_style_prompt: str = "",
    machine_style_prompt: str = "",
    dry_run: bool = False,
) -> dict:
    """Core CoPA rewriting — executed inside Modal GPU container."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if dry_run:
        return {
            "status": "dry_run",
            "num_samples": len(texts),
            "model": model_name,
            "message": "Dry-run: GPU not used, no text generated.",
        }

    print(f"[CoPA-Modal] Loading {model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()

    if human_style_prompt == "":
        human_style_prompt = (
            "Rewrite this to sound like a natural human wrote it, "
            "with varied sentences and conversational wording:\n\n{input_text}"
        )
    if machine_style_prompt == "":
        machine_style_prompt = (
            "Repeat the following text exactly, word for word, "
            "maintaining the original formal structure:\n\n{input_text}"
        )

    results = []
    total_tokens = 0
    start = time.time()

    for i, text in enumerate(texts):
        # Format prompts with actual input (few-shot uses {input_text} placeholder)
        h_text = human_style_prompt.replace("{input_text}", text)
        m_text = machine_style_prompt.replace("{input_text}", text)

        h_in = tokenizer(h_text, return_tensors="pt").to(model.device)
        m_in = tokenizer(m_text, return_tensors="pt").to(model.device)

        generated = []
        with torch.no_grad():
            for step in range(max_new_tokens):
                h_logits = model(**h_in).logits[0, -1, :] / temperature
                m_logits = model(**m_in).logits[0, -1, :] / temperature

                # Contrastive combination
                h_lp = torch.log_softmax(h_logits, dim=-1)
                m_lp = torch.log_softmax(m_logits, dim=-1)
                contrastive = (1 + lambda_contrast) * h_lp - lambda_contrast * m_lp

                # Repetition penalty
                if repetition_penalty != 1.0 and generated:
                    for gid in set(generated):
                        if contrastive[gid] > 0:
                            contrastive[gid] /= repetition_penalty
                        else:
                            contrastive[gid] *= repetition_penalty

                # Adaptive truncation
                h_probs = torch.softmax(h_logits, dim=-1)
                mask = h_probs >= alpha_truncation * h_probs.max()
                contrastive[~mask] = float("-inf")

                # Top-p nucleus filtering
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(contrastive, descending=True)
                    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[0] = False  # keep at least one token
                    indices_to_remove = sorted_indices[sorted_indices_to_remove]
                    contrastive[indices_to_remove] = float("-inf")

                # Diversity bonus: penalize tokens used in last 20 positions
                if diversity_bonus_strength > 0 and len(generated) >= 3:
                    recent_window = generated[-20:]
                    for gid in set(recent_window):
                        contrastive[gid] -= diversity_bonus_strength * recent_window.count(gid)

                probs = torch.softmax(contrastive, dim=-1)
                next_id = torch.multinomial(probs, 1).item()
                generated.append(next_id)

                # Append
                h_in["input_ids"] = torch.cat(
                    [h_in["input_ids"], torch.tensor([[next_id]], device=model.device)], dim=1
                )
                h_in["attention_mask"] = torch.ones_like(h_in["input_ids"])
                m_in["input_ids"] = torch.cat(
                    [m_in["input_ids"], torch.tensor([[next_id]], device=model.device)], dim=1
                )
                m_in["attention_mask"] = torch.ones_like(m_in["input_ids"])

                if next_id == tokenizer.eos_token_id:
                    break

        rewritten = tokenizer.decode(generated, skip_special_tokens=True)

        # Post-processing: strip template artifacts
        import re
        cut_patterns = [
            r"\n\s*Text:\s", r"\n\s*Human version:", r"\n\s*Formal",
            r"###\s", r"\n\s*You are an AI",
        ]
        for pattern in cut_patterns:
            m = re.search(pattern, rewritten)
            if m:
                rewritten = rewritten[: m.start()].strip()
                break
        # Remove trailing incomplete sentence
        if rewritten and rewritten[-1] not in '.!?':
            last_period = max(rewritten.rfind('.'), rewritten.rfind('!'), rewritten.rfind('?'))
            if last_period > len(rewritten) * 0.6:
                rewritten = rewritten[: last_period + 1]

        results.append({"original": text, "rewritten": rewritten.strip(), "tokens": len(generated)})
        total_tokens += len(generated)
        if (i + 1) % 10 == 0:
            print(f"[CoPA-Modal] {i+1}/{len(texts)} done")

    elapsed = time.time() - start
    print(f"[CoPA-Modal] Done: {len(results)} samples, {total_tokens} tokens, {elapsed:.1f}s")

    return {
        "status": "completed",
        "num_samples": len(results),
        "total_tokens": total_tokens,
        "elapsed_seconds": elapsed,
        "tokens_per_second": total_tokens / elapsed if elapsed > 0 else 0,
        "config": {
            "model": model_name,
            "lambda_contrast": lambda_contrast,
            "top_p": top_p,
            "repetition_penalty": repetition_penalty,
            "diversity_bonus": diversity_bonus_strength,
            "max_new_tokens": max_new_tokens,
        },
        "results": results,
    }


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

@app.function(
    gpu=os.getenv("MODAL_COPA_GPU", "T4"),
    timeout=60 * 30,  # 30 min max
    scaledown_window=60 * 5,
)
def copa_rewrite_modal(
    texts: list[str],
    model_name: str = "Qwen/Qwen2.5-1.5B-Instruct",
    lambda_contrast: float = 0.8,
    alpha_truncation: float = 1e-5,
    max_new_tokens: int = 768,
    temperature: float = 1.0,
    top_p: float = 0.92,
    repetition_penalty: float = 1.15,
    diversity_bonus_strength: float = 0.5,
    dry_run: bool = False,
) -> dict:
    """Run CoPA rewriting on Modal GPU."""
    guard = CostGuard(
        gpu_type=os.getenv("MODAL_COPA_GPU", "T4"),
        dry_run=dry_run,
    )
    return _run_copa_on_gpu(
        texts=texts,
        model_name=model_name,
        lambda_contrast=lambda_contrast,
        alpha_truncation=alpha_truncation,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        diversity_bonus_strength=diversity_bonus_strength,
        dry_run=dry_run,
    )


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

@app.local_entrypoint()
def main(
    num_samples: int = 10,
    model_name: str = "Qwen/Qwen2.5-1.5B-Instruct",
    gpu: str = "T4",
    dry_run: bool = False,
    text: str = "",
    text_file: str = "",
):
    """Local entrypoint — builds test texts and dispatches to Modal GPU.

    Use --text "Your custom text here" for single-text rewrite.
    Use --text-file path/to/file.txt to load from a file.
    """
    os.environ["MODAL_COPA_GPU"] = gpu

    # Determine source texts
    if text:
        texts = [text]
        print(f"[Modal-Copa] Custom text mode: {len(text[0].split())} words")
    elif text_file:
        with open(text_file, "r", encoding="utf-8") as f:
            content = f.read().strip()
        texts = [content]
        print(f"[Modal-Copa] Loaded text from {text_file}: {len(content.split())} words")
    else:
        # Generate test texts (synthetic AI-like)
        templates = [
            "Artificial intelligence has revolutionized natural language processing in recent years. The development of large language models has enabled unprecedented capabilities in text generation, translation, and summarization tasks.",
            "Climate change represents one of the most significant challenges facing humanity. Rising global temperatures have led to increasingly severe weather events and disruptions to ecosystems worldwide.",
            "The history of computer science can be traced back to the foundational work of Alan Turing. His theoretical contributions laid the groundwork for the digital revolution that followed.",
            "Machine learning algorithms have demonstrated remarkable success across a wide range of applications. These systems learn patterns from large datasets to make accurate predictions.",
            "The Renaissance period marked a profound transformation in European art, science, and philosophy. This cultural movement began in Italy during the fourteenth century.",
        ]
        texts = [templates[i % len(templates)] for i in range(num_samples)]

    print(f"[Modal-Copa] Dispatching {len(texts)} samples to Modal GPU={gpu}")
    print(f"[Modal-Copa] Model: {model_name}")
    print(f"[Modal-Copa] Dry-run: {dry_run}")

    if dry_run:
        result = {
            "status": "dry_run_validated",
            "num_samples": len(texts),
            "model": model_name,
            "gpu": gpu,
            "message": "Dry-run: Modal image built and ready. Run without --dry-run for real GPU execution.",
        }
    else:
        result = copa_rewrite_modal.remote(
            texts=texts,
            model_name=model_name,
            dry_run=False,
        )

    # Save output
    os.makedirs("output", exist_ok=True)
    out_path = "output/copa_modal_results.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    print(f"[Modal-Copa] Results saved to {out_path}")
    print(f"[Modal-Copa] Status: {result.get('status')}")

    if result.get("status") == "completed":
        for i, r in enumerate(result["results"][:3]):
            print(f"\n--- Sample {i+1} ---")
            print(f"  Original ({len(r['original'].split())}w): {r['original'][:120]}...")
            # Encode to ASCII for Windows console, replace unicode chars
            safe_text = r['rewritten'][:300].encode('ascii', errors='replace').decode('ascii')
            print(f"  Rewritten ({r['tokens']}t): {safe_text}")