#!/usr/bin/env python
"""
AI Text Rewriting — Evasion Detection Pipeline

Rewrites AI-generated text into natural, human-like prose that evades
AI text detectors (Fast-DetectGPT, Binoculars, GPTZero, Pangram).

Usage:
  python rewrite_text.py "Your AI-generated text here"
  python rewrite_text.py --file data/bitcoin_text.txt
  python rewrite_text.py --file input.txt --verify --stats
  python rewrite_text.py --help

How it works:
  1. Sends text to Modal GPU (T4, ~0.60/h)
  2. Qwen2.5-1.5B-Instruct rewrites with optimized sampling
  3. Verifies: coherence, length similarity, no artifacts, no hallucinations
  4. Computes dispersion statistics (TTR, burstiness, word freq variance)
  5. Saves result + uploads to HuggingFace

Requirements: Modal CLI installed and configured
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
import time
from pathlib import Path

SCRIPT_DIR = Path(__file__).parent
MODAL_APP = SCRIPT_DIR / "src" / "modal_app_rewrite.py"


def run_rewrite(text: str, gpu: str = "T4", verify: bool = True,
                dry_run: bool = False) -> dict:
    """Dispatch rewrite to Modal GPU."""
    print(f"\n[Rewrite] Dispatching to Modal {gpu} GPU...")
    print(f"[Rewrite] Input: {len(text.split())} words, {len(text)} chars")

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".txt", delete=False, encoding="utf-8"
    ) as f:
        f.write(text)
        temp_path = f.name

    try:
        cmd = [
            "modal", "run", "-q",
            str(MODAL_APP),
            "--text-file", temp_path,
            "--gpu", gpu,
        ]
        if verify:
            cmd.append("--verify")
        if dry_run:
            cmd.append("--dry-run")

        start = time.time()
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            cwd=str(SCRIPT_DIR),
            timeout=600,
            env={**os.environ, "PYTHONIOENCODING": "utf-8"},
        )

        elapsed = time.time() - start

        out_file = SCRIPT_DIR / "output" / "rewrite_result.json"
        if out_file.exists():
            with open(out_file, "r", encoding="utf-8") as f:
                data = json.load(f)
            if data.get("status") == "completed":
                data["_dispatch_time_s"] = round(elapsed, 1)
                return data

        if result.returncode != 0 and "UnicodeEncodeError" not in result.stderr:
            return {"status": "error", "message": result.stderr[-500:]}

        return {"status": "error", "message": "No output file"}

    except subprocess.TimeoutExpired:
        return {"status": "error", "message": "Timeout (10 min)"}
    finally:
        os.unlink(temp_path)


def verify_rewrite(original: str, rewritten: str) -> dict:
    """Verify rewrite quality."""
    orig_w = len(original.split())
    rew_w = len(rewritten.split())
    ratio = rew_w / max(orig_w, 1)

    issues = []
    ok = []

    # 1. Length
    if ratio < 0.4:
        issues.append(f"Too short: {rew_w}w vs {orig_w}w (ratio={ratio:.2f})")
    elif ratio > 2.5:
        issues.append(f"Too long: {rew_w}w vs {orig_w}w (ratio={ratio:.2f})")
    else:
        ok.append(f"Length: {orig_w}w -> {rew_w}w (ratio={ratio:.2f})")

    # 2. Artifacts
    artifacts = ["###", "Here's a", "Let me know", "Feel free", "I hope this",
                 "In conclusion", "To summarize", "Note:", "Please note"]
    found_artifacts = [a for a in artifacts if a.lower() in rewritten.lower()]
    if found_artifacts:
        issues.append(f"Artifacts: {found_artifacts}")
    else:
        ok.append("No artifacts detected")

    # 3. Key facts preserved
    orig_nums = set(re.findall(r'\b\d+\b', original))
    rew_nums = set(re.findall(r'\b\d+\b', rewritten))
    if orig_nums:
        missing = orig_nums - rew_nums
        if len(missing) / len(orig_nums) > 0.5:
            issues.append(f"Missing numbers: {missing}")
        else:
            ok.append(f"Numbers preserved: {len(orig_nums - missing)}/{len(orig_nums)}")

    # 4. Proper nouns preserved
    orig_proper = set(re.findall(r'\b[A-Z][a-z]{2,}\b', original))
    rew_lower = rewritten.lower()
    missing_proper = [p for p in orig_proper if p.lower() not in rew_lower]
    if len(missing_proper) > len(orig_proper) * 0.5:
        issues.append(f"Missing key terms: {missing_proper[:5]}")
    else:
        ok.append(f"Key terms preserved: {len(orig_proper) - len(missing_proper)}/{len(orig_proper)}")

    # 5. Repetition
    sentences = [s.strip() for s in re.split(r'[.!?]+', rewritten) if len(s.strip()) > 10]
    if len(sentences) >= 3:
        unique = len(set(s[:30].lower() for s in sentences))
        if unique < len(sentences) * 0.5:
            issues.append(f"Repetitive: {unique} unique / {len(sentences)} sentences")
        else:
            ok.append(f"Varied sentences: {unique} unique / {len(sentences)}")

    # 6. No new elements
    orig_set = set(original.lower().split())
    rew_set = set(rewritten.lower().split())
    new_words = rew_set - orig_set
    if len(new_words) > len(rew_set) * 0.7:
        issues.append(f"Too many new words: {len(new_words)}/{len(rew_set)}")
    else:
        ok.append(f"Vocabulary overlap: {len(rew_set - new_words)}/{len(rew_set)} original")

    return {
        "passed": len(issues) == 0,
        "ok": ok,
        "issues": issues,
        "length_ratio": round(ratio, 2),
        "original_words": orig_w,
        "rewritten_words": rew_w,
    }


def compute_stats(original: str, rewritten: str) -> dict:
    """Compute dispersion statistics."""
    sys.path.insert(0, str(SCRIPT_DIR / "src"))
    from eval_statistical import compute_stats as cs, compute_dispersion_score as cds

    os_, rs = cs(original), cs(rewritten)
    od, rd = cds(os_), cds(rs)

    return {
        "original": {"words": os_.num_words, "ttr": round(os_.type_token_ratio, 3),
                     "sent_cv": round(os_.sentence_len_cv, 3),
                     "word_freq_std": round(os_.std_word_freq, 2),
                     "flesch": round(os_.readability_flesch, 0),
                     "human_likeness": od["overall_human_likeness"]},
        "rewritten": {"words": rs.num_words, "ttr": round(rs.type_token_ratio, 3),
                      "sent_cv": round(rs.sentence_len_cv, 3),
                      "word_freq_std": round(rs.std_word_freq, 2),
                      "flesch": round(rs.readability_flesch, 0),
                      "human_likeness": rd["overall_human_likeness"]},
        "deltas": {k: round(rd[k] - od[k], 3) for k in od},
    }


def main():
    parser = argparse.ArgumentParser(
        description="Rewrite AI text to evade detectors (Modal GPU)",
        epilog="Example: python rewrite_text.py 'Your AI text here' --verify --stats"
    )
    parser.add_argument("text", nargs="?", help="Text to rewrite (or use --file)")
    parser.add_argument("--file", "-f", help="File containing text")
    parser.add_argument("--output", "-o", default="output/rewrite_result.json")
    parser.add_argument("--gpu", default="T4")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--no-verify", action="store_true")
    parser.add_argument("--stats", action="store_true")
    parser.add_argument("--upload", action="store_true", help="Upload to HF after rewrite")
    args = parser.parse_args()

    # Input
    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            text = f.read().strip()
    elif args.text:
        text = args.text
    else:
        print("ERROR: Provide text or use --file")
        sys.exit(1)

    if not text.strip():
        print("ERROR: Empty text")
        sys.exit(1)

    print("=" * 60)
    print(" AI Text Rewriting — Evasion Detection")
    print("=" * 60)
    print(f"\n[Original] {len(text.split())} words, {len(text)} chars:")
    print(text[:300] + ("..." if len(text) > 300 else ""))

    # Run
    result = run_rewrite(text, gpu=args.gpu, verify=not args.no_verify,
                         dry_run=args.dry_run)

    if result.get("status") != "completed":
        print(f"\nERROR: {result.get('message', 'Unknown error')}")
        sys.exit(1)

    r = result["results"][0]
    rewritten = r["rewritten"]

    print(f"\n[Rewritten] {len(rewritten.split())} words, {r.get('tokens', '?')} tokens, "
          f"{r.get('elapsed_seconds', '?')}s:")
    print(rewritten)

    # Verification
    if not args.no_verify:
        print(f"\n{'=' * 60}")
        print(" QUALITY VERIFICATION")
        print("=" * 60)
        v = r.get("verification") or verify_rewrite(text, rewritten)

        for check in v.get("ok", []):
            print(f"  OK   {check}")
        for issue in v.get("issues", []):
            print(f"  FAIL {issue}")

        status = "ALL CHECKS PASSED" if v["passed"] else "SOME CHECKS FAILED"
        print(f"\n  => {status}")
        print(f"  => Length: {v['original_words']}w -> {v['rewritten_words']}w "
              f"(ratio: {v['length_ratio']:.2f})")

    # Statistics
    if args.stats:
        print(f"\n{'=' * 60}")
        print(" DISPERSION STATISTICS")
        print("=" * 60)
        s = compute_stats(text, rewritten)
        print(f"  {'Metric':<25} {'Original':>10} {'Rewritten':>10} {'Delta':>10}")
        print(f"  {'-'*55}")
        for metric in ["words", "ttr", "sent_cv", "word_freq_std", "flesch", "human_likeness"]:
            o_val = s["original"][metric]
            r_val = s["rewritten"][metric]
            d_val = s["deltas"][metric]
            print(f"  {metric:<25} {str(o_val):>10} {str(r_val):>10} {d_val:>+10}")

    # Save
    output = {
        "original": text,
        "rewritten": rewritten,
        "model": r.get("model", "Qwen/Qwen2.5-1.5B-Instruct"),
        "method": "direct-generation-optimized-sampling",
        "config": result.get("config", {}),
        "tokens": r.get("tokens", 0),
        "elapsed_seconds": r.get("elapsed_seconds", 0),
    }
    if not args.no_verify:
        output["verification"] = v
    if args.stats:
        output["statistics"] = s

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    print(f"\n[Save] {args.output}")
    print(f"[Save] View: https://huggingface.co/simonlesaumon/evasion-detection-artifacts")

    # Upload
    if args.upload:
        import subprocess as sp
        sp.run(["hf", "upload", "simonlesaumon/evasion-detection-artifacts",
                args.output, f"results/rewrite_{int(time.time())}.json"],
               cwd=str(SCRIPT_DIR))

    return 0 if (v.get("passed", True) if not args.no_verify else True) else 1


if __name__ == "__main__":
    sys.exit(main())