"""
AI vs Human Text Detector — Inference API

Wraps the saved RoBERTa classifier (ai-detector-model-v2) in a small FastAPI
service so HumanPen (or any browser-based tool) can call it over HTTP, since
the model can't run directly in-browser like Groq/Anthropic API calls do.

Endpoints:
  GET  /            -> health check
  POST /detect       -> score a single piece of text (paragraph-level recommended)
  POST /detect_batch -> score multiple pieces of text in one call (more efficient
                         for HumanPen's "scan whole document" use case)

Designed to run on Hugging Face Spaces (Docker SDK) but works anywhere that
can run a Python container exposing port 7860.
"""

import os
from typing import List

import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import RobertaTokenizer, RobertaForSequenceClassification

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

# On Hugging Face Spaces, you'll upload the model files into a folder named
# "model" alongside this app.py (see DEPLOY.md for the exact layout).
MODEL_PATH = os.environ.get("MODEL_PATH", "./model")

# Matches what the training notebook used — texts longer than this were
# truncated during training, so keep inference consistent with that.
MAX_CHARS = 2000
MAX_TOKENS = 512

# ---------------------------------------------------------------------------
# Load model once at startup (NOT per-request — this is the expensive part)
# ---------------------------------------------------------------------------

print(f"Loading tokenizer from {MODEL_PATH} ...")
tokenizer = RobertaTokenizer.from_pretrained(MODEL_PATH)

print(f"Loading model from {MODEL_PATH} ...")
model = RobertaForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(f"Model loaded. Using device: {device}")


def score_text(text: str) -> dict:
    """Run a single piece of text through the model and return a clean result.

    Returns:
        {
          "label": "AI" | "Human",
          "confidence": float (0-1, confidence in the predicted label),
          "ai_probability": float (0-1, raw probability of the AI class,
                                    useful for a continuous heatmap score
                                    rather than just a binary verdict),
          "truncated": bool (whether input was cut to MAX_CHARS)
        }
    """
    truncated = len(text) > MAX_CHARS
    text_for_model = text[:MAX_CHARS]

    inputs = tokenizer(
        text_for_model,
        truncation=True,
        padding=True,
        max_length=MAX_TOKENS,
        return_tensors="pt",
    ).to(device)

    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1)[0]
    ai_probability = float(probs[1])  # label 1 = AI, per training notebook
    pred_label = int(torch.argmax(probs))
    confidence = float(probs[pred_label])

    return {
        "label": "AI" if pred_label == 1 else "Human",
        "confidence": round(confidence, 4),
        "ai_probability": round(ai_probability, 4),
        "truncated": truncated,
    }


# ---------------------------------------------------------------------------
# API schema
# ---------------------------------------------------------------------------

class DetectRequest(BaseModel):
    text: str


class DetectBatchRequest(BaseModel):
    texts: List[str]


class DetectResponse(BaseModel):
    label: str
    confidence: float
    ai_probability: float
    truncated: bool


class DebugRequest(BaseModel):
    text: str


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(title="AI vs Human Text Detector API")

# Allow HumanPen (running as a local HTML file or hosted elsewhere) to call
# this from the browser. Restrict allow_origins in production if you want
# to lock this down to a specific domain instead of "*".
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)


@app.get("/")
def health_check():
    return {"status": "ok", "model": "ai-detector-model-v2", "device": str(device)}


@app.post("/detect", response_model=DetectResponse)
def detect(req: DetectRequest):
    if not req.text or not req.text.strip():
        raise HTTPException(status_code=400, detail="Text must not be empty.")
    return score_text(req.text)


@app.post("/detect_batch")
def detect_batch(req: DetectBatchRequest):
    if not req.texts:
        raise HTTPException(status_code=400, detail="texts list must not be empty.")
    results = []
    for text in req.texts:
        if not text or not text.strip():
            results.append({"label": "Human", "confidence": 0.0, "ai_probability": 0.0, "truncated": False, "skipped": True})
            continue
        results.append(score_text(text))
    return {"results": results}


@app.post("/debug")
def debug(req: DebugRequest):
    """TEMPORARY diagnostic endpoint — remove once the deployment bug is found.
    Returns raw logits, token IDs, and model/tokenizer fingerprints so we can
    see exactly what's happening inside the container for a given input.
    """
    text = req.text[:MAX_CHARS]
    inputs = tokenizer(
        text, truncation=True, padding=True, max_length=MAX_TOKENS, return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1)[0]

    first_param = next(model.parameters())
    weight_fingerprint = float(first_param.flatten()[0].item())

    classifier_fingerprint = None
    for name, param in model.named_parameters():
        if "classifier" in name:
            classifier_fingerprint = {
                "param_name": name,
                "shape": list(param.shape),
                "first_5_values": param.flatten()[:5].tolist(),
            }
            break

    import transformers as _tf
    import torch as _torch

    return {
        "input_text_received": text[:100],
        "input_text_length": len(text),
        "token_ids_first_10": inputs["input_ids"][0][:10].tolist(),
        "token_ids_last_10": inputs["input_ids"][0][-10:].tolist(),
        "num_tokens": int(inputs["input_ids"].shape[1]),
        "raw_logits": logits[0].tolist(),
        "softmax_probs": probs.tolist(),
        "model_weight_fingerprint": weight_fingerprint,
        "model_training_mode": model.training,
        "classifier_head_fingerprint": classifier_fingerprint,
        "transformers_version": _tf.__version__,
        "torch_version": _torch.__version__,
    }