import torch
import math
import re
import statistics
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

app = FastAPI(title="Perplexity + Burstiness API — UltimateEdge")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Load model at startup ───────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading GPT-2 medium on {device}...")
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2-medium")
model = GPT2LMHeadModel.from_pretrained("gpt2-medium").to(device)
model.eval()
print("Model ready.")


# ── Request schema ──────────────────────────────────────────
class TextRequest(BaseModel):
    text: str


# ── Core functions ──────────────────────────────────────────
def compute_perplexity(text: str) -> dict:
    encodings = tokenizer(text, return_tensors="pt")
    input_ids = encodings.input_ids.to(device)
    seq_len = input_ids.size(1)
    if seq_len == 0:
        return {"error": "Empty text."}

    stride, max_length = 512, 1024
    nlls, prev_end = [], 0

    for begin_loc in range(0, seq_len, stride):
        end_loc = min(begin_loc + max_length, seq_len)
        trg_len = end_loc - prev_end
        chunk = input_ids[:, begin_loc:end_loc]
        labels = chunk.clone()
        labels[:, :-trg_len] = -100
        with torch.no_grad():
            loss = model(chunk, labels=labels).loss
        nlls.append(loss * trg_len)
        prev_end = end_loc
        if end_loc == seq_len:
            break

    mean_nll = torch.stack(nlls).sum() / seq_len
    ppl = torch.exp(mean_nll).item()

    if ppl < 20:    ppl_label = "Very Low — AI-like"
    elif ppl < 50:  ppl_label = "Low — fluent prose"
    elif ppl < 100: ppl_label = "Moderate — human range"
    elif ppl < 200: ppl_label = "High — complex/varied"
    else:           ppl_label = "Very High — noisy/technical"

    return {
        "perplexity": round(ppl, 2),
        "bits_per_token": round(mean_nll.item() / math.log(2), 4),
        "n_tokens": int(seq_len),
        "interpretation": ppl_label,
    }


def get_sentence_ppls(text: str) -> list:
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    results = []
    for sent in sentences:
        sent = sent.strip()
        if len(sent.split()) < 3:
            continue
        try:
            r = compute_perplexity(sent)
            if "error" not in r:
                results.append({"sentence": sent, "ppl": r["perplexity"], "tokens": r["n_tokens"]})
        except Exception:
            continue
    return results


def compute_burstiness(ppl_scores: list) -> dict:
    if len(ppl_scores) < 2:
        return {"error": "Need at least 2 sentences."}
    mean = statistics.mean(ppl_scores)
    std = statistics.stdev(ppl_scores)
    burstiness = (std - mean) / (std + mean) if (std + mean) != 0 else 0

    if burstiness > 0.2:    b_label = "Very High — strong human signature"
    elif burstiness > 0.05: b_label = "High — likely human writing"
    elif burstiness > -0.1: b_label = "Moderate — mixed signals"
    elif burstiness > -0.3: b_label = "Low — leans AI-generated"
    else:                   b_label = "Very Low — strong AI signature"

    return {
        "burstiness_score": round(burstiness, 4),
        "std_dev": round(std, 2),
        "mean_sentence_ppl": round(mean, 2),
        "coefficient_of_variation_pct": round((std / mean * 100) if mean else 0, 1),
        "interpretation": b_label,
    }


def compute_verdict(ppl: float, burstiness: float) -> dict:
    if ppl < 20:      ppl_score = 5
    elif ppl < 40:    ppl_score = 25
    elif ppl < 60:    ppl_score = 55
    elif ppl < 100:   ppl_score = 75
    elif ppl < 200:   ppl_score = 85
    else:             ppl_score = 60

    if burstiness > 0.3:    b_score = 95
    elif burstiness > 0.1:  b_score = 80
    elif burstiness > 0.0:  b_score = 60
    elif burstiness > -0.1: b_score = 40
    elif burstiness > -0.3: b_score = 20
    else:                   b_score = 5

    human_score = round((ppl_score * 0.45) + (b_score * 0.55))

    if human_score >= 75:   verdict, detail = "✅ Likely Human-Written", "Both metrics suggest natural, varied writing."
    elif human_score >= 55: verdict, detail = "⚠️ Mixed — Possibly AI-Assisted", "Some human patterns but uniformity in places."
    elif human_score >= 35: verdict, detail = "🤖 Leans AI-Generated", "Low burstiness and predictable structure."
    else:                   verdict, detail = "🤖 Very Likely AI-Generated", "Strong AI signature across both metrics."

    return {
        "verdict": verdict,
        "human_score": f"{human_score}/100",
        "detail": detail,
        "ppl_contribution": f"{ppl_score}/100",
        "burstiness_contribution": f"{b_score}/100",
    }


# ── Routes ──────────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "name": "Perplexity + Burstiness API — UltimateEdge",
        "usage": "POST /analyse with JSON body: {\"text\": \"your article here\"}",
        "endpoints": ["/analyse", "/health"]
    }

@app.get("/health")
def health():
    return {"status": "ok", "device": device}

@app.post("/analyse")
def analyse(req: TextRequest):
    text = req.text.strip()
    if not text:
        return {"error": "Empty text provided."}

    overall = compute_perplexity(text)
    if "error" in overall:
        return overall

    sent_results = get_sentence_ppls(text)
    ppl_scores = [s["ppl"] for s in sent_results]
    burst = compute_burstiness(ppl_scores) if len(ppl_scores) >= 2 else {"error": "Not enough sentences."}
    b_val = burst.get("burstiness_score", 0)
    verdict = compute_verdict(overall["perplexity"], b_val)

    return {
        "perplexity": overall,
        "burstiness": burst,
        "verdict": verdict,
        "sentences": sorted(sent_results, key=lambda x: x["ppl"]),
    }