"""
UnifIDE AI Detector — v3 (Binary UniXcoder + Platt Calibration)
Model: microsoft/unixcoder-base fine-tuned as binary classifier
  human (0) | ai (1)
AI% = sigmoid(logit) calibrated via Platt scaling
"""

import os
import json
import pickle
import contextlib
import numpy as np
import torch
from torch import nn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

# ── Config ─────────────────────────────────────────────────────────────────────
MODEL_REPO        = os.getenv("AI_MODEL_REPO", "ss123qwd/unifide-ai-detector-v3")
HF_TOKEN          = os.getenv("HF_TOKEN", None)
OPERATE_THRESHOLD = float(os.getenv("SUBMIT_THRESHOLD", "51.4"))
MIN_CODE_LEN      = 150
MAX_LENGTH        = 512  # must match training

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[AI-DETECTOR] Device   : {DEVICE}")
print(f"[AI-DETECTOR] Model    : {MODEL_REPO}")
print(f"[AI-DETECTOR] Threshold: {OPERATE_THRESHOLD}%")


# ── Model architecture — must match training exactly ───────────────────────────
class UniXcoderBinaryClassifier(nn.Module):
    """
    Binary sigmoid classifier.
    Output: scalar logit. sigmoid(logit) = P(AI).
    AI% = P(AI) * 100.
    """
    def __init__(self, model_name, dropout=0.1):
        super().__init__()
        self.encoder    = AutoModel.from_pretrained(model_name)
        self.dropout    = nn.Dropout(dropout)
        self.classifier = nn.Linear(self.encoder.config.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        out     = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_emb = out.last_hidden_state[:, 0, :]
        return self.classifier(self.dropout(cls_emb)).squeeze(-1)


# ── Download helpers ───────────────────────────────────────────────────────────
def download(filename):
    return hf_hub_download(
        repo_id=MODEL_REPO,
        filename=filename,
        token=HF_TOKEN,
    )


# ── Load calibrator ────────────────────────────────────────────────────────────
print("[AI-DETECTOR] Loading calibrator...")
cal_path   = download("calibrator.pkl")
with open(cal_path, "rb") as f:
    calibrator = pickle.load(f)
print("[AI-DETECTOR] Calibrator loaded.")

# ── Load threshold ─────────────────────────────────────────────────────────────
try:
    threshold_path = download("threshold.json")
    with open(threshold_path) as f:
        saved = json.load(f)
    TUNED_THRESHOLD = float(saved.get("submit_threshold", 51.4))
    print(f"[AI-DETECTOR] Tuned threshold from file : {TUNED_THRESHOLD}%")
    print(f"[AI-DETECTOR] Operating threshold (used): {OPERATE_THRESHOLD}%")
except Exception as e:
    print(f"[AI-DETECTOR] WARN: could not load threshold.json: {e}")
    TUNED_THRESHOLD = 51.4

# ── Load tokenizer ─────────────────────────────────────────────────────────────
print("[AI-DETECTOR] Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, token=HF_TOKEN)
print("[AI-DETECTOR] Tokenizer loaded.")

# ── Load model weights ─────────────────────────────────────────────────────────
print("[AI-DETECTOR] Loading model weights (~30s)...")
model = UniXcoderBinaryClassifier(
    "microsoft/unixcoder-base", dropout=0.1
).to(DEVICE)

weights_path = download("best_model.pt")
model.load_state_dict(
    torch.load(weights_path, map_location=DEVICE)
)
model.eval()
print("[AI-DETECTOR] Model ready.")


# ── Inference ──────────────────────────────────────────────────────────────────
@torch.no_grad()
def predict(code: str, language: str) -> dict:
    """
    Run binary inference and return result dict.

    Actions:
      pass           — AI% <= threshold, allow submission
      flag_ai        — AI% > threshold AND model is confident
      pending_review — AI% > threshold BUT model is uncertain (borderline)
    """
    if len(code.strip()) < MIN_CODE_LEN:
        return {
            "ai_probability":    0.0,
            "human_probability": 1.0,
            "ai_percent":        0.0,
            "label":             "Too Short",
            "action":            "pending_review",
            "confidence":        None,
            "variant":           None,
            "breakdown":         None,
            "warning": f"Code too short (< {MIN_CODE_LEN} chars). Manual review recommended.",
        }

    # Prepend language token — same as training
    lang_safe       = language.strip().replace(" ", "_").replace("#", "sharp")
    code_with_token = f"<lang:{lang_safe}> {code}"

    enc = tokenizer(
        code_with_token,
        max_length=MAX_LENGTH,
        truncation=True,
        padding=True,
        return_tensors="pt",
    )

    # Device-aware inference — no hardcoded 'cuda'
    with contextlib.nullcontext():
        logit = model(
            enc["input_ids"].to(DEVICE),
            enc["attention_mask"].to(DEVICE),
        ).cpu().float().item()

    # Raw sigmoid probability
    raw_prob = float(torch.sigmoid(torch.tensor(logit)).item())

    # Calibrated AI% via Platt scaling
    cal_prob = float(calibrator.predict_proba([[logit]])[:, 1][0])
    ai_pct   = round(cal_prob * 100, 2)

    # Confidence based on distance from decision boundary
    dist_from_half = abs(raw_prob - 0.5)
    if dist_from_half >= 0.25:
        confidence = "high"
    elif dist_from_half >= 0.10:
        confidence = "medium"
    else:
        confidence = "low"

    # Action logic
    above_threshold = ai_pct > OPERATE_THRESHOLD
    if not above_threshold:
        action = "pass"
    elif confidence == "high":
        action = "flag_ai"
    else:
        action = "pending_review"

    human_prob = 1.0 - cal_prob

    return {
        # Backward-compatible fields
        "ai_probability":    round(cal_prob, 4),
        "human_probability": round(human_prob, 4),
        "ai_percent":        ai_pct,
        "label":             "AI-Generated" if action == "flag_ai" else "Human-Written",

        # New fields
        "action":           action,
        "confidence":       confidence,
        "variant":          None,   # binary model has no variant — kept for API compat
        "breakdown":        {
            "human": round(human_prob * 100, 2),
            "ai":    round(cal_prob * 100, 2),
        },
        "threshold_used":   OPERATE_THRESHOLD,
        "tuned_threshold":  TUNED_THRESHOLD,
    }


# ── FastAPI app ────────────────────────────────────────────────────────────────
app = FastAPI(title="UnifIDE AI Detector v3")

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


class DetectRequest(BaseModel):
    code: str
    language: str


@app.get("/health")
def health():
    return {
        "ok":        True,
        "model":     MODEL_REPO,
        "version":   "v3",
        "threshold": OPERATE_THRESHOLD,
    }


@app.post("/detect-code")
@app.post("/detect-code/")
def detect_code(payload: DetectRequest):
    code = payload.code.strip()
    if not code:
        raise HTTPException(status_code=400, detail="Code is empty.")
    try:
        result = predict(code, payload.language)
    except Exception as e:
        print("[AI-DETECTOR] Prediction error:", e)
        raise HTTPException(status_code=500, detail="Model prediction failed.")
    return result