import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import re
import pandas as pd
import gradio as gr

# -----------------------------
# MODEL INITIALIZATION
# -----------------------------
MODEL_NAME = "fakespot-ai/roberta-base-ai-text-detection-v1"
tokenizer = None
model = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def get_model():
    global tokenizer, model
    if model is None:
        print(f"Loading model: {MODEL_NAME} on {device}")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        dtype = torch.float32
        if device.type == "cuda" and torch.cuda.is_bf16_supported():
            dtype = torch.bfloat16
        model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_NAME, torch_dtype=dtype
        ).to(device).eval()
    return tokenizer, model

# UPDATED THRESHOLD: Only 81% and above is flagged as AI
THRESHOLD = 0.81 

# -----------------------------
# PROTECT STRUCTURE
# -----------------------------
ABBR = ["e.g", "i.e", "mr", "mrs", "ms", "dr", "prof", "vs", "etc", "fig", "al", "jr", "sr", "st", "inc", "ltd", "u.s", "u.k"]
ABBR_REGEX = re.compile(r"\b(" + "|".join(map(re.escape, ABBR)) + r")\.", re.IGNORECASE)

def _protect(text):
    text = text.replace("...", "⟨ELLIPSIS⟩")
    text = re.sub(r"(?<=\d)\.(?=\d)", "⟨DECIMAL⟩", text)
    text = ABBR_REGEX.sub(r"\1⟨ABBRDOT⟩", text)
    return text

def _restore(text):
    return text.replace("⟨ABBRDOT⟩", ".").replace("⟨DECIMAL⟩", ".").replace("⟨ELLIPSIS⟩", "...")

def split_preserving_structure(text):
    blocks = re.split(r"(\n+)", text)
    final_blocks = []
    for block in blocks:
        if block.startswith("\n"):
            final_blocks.append(block)
        else:
            protected = _protect(block)
            parts = re.split(r"([.?!])(\s+)", protected)
            for i in range(0, len(parts), 3):
                sentence = parts[i]
                punct = parts[i+1] if i+1 < len(parts) else ""
                space = parts[i+2] if i+2 < len(parts) else ""
                if sentence.strip():
                    final_blocks.append(_restore(sentence + punct))
                if space:
                    final_blocks.append(space)
    return final_blocks

# -----------------------------
# ANALYSIS
# -----------------------------
@torch.inference_mode()
def analyze(text):
    text = text.strip()
    if not text:
        return "—", "—", "<em>Please enter text...</em>", None
        
    word_count = len(text.split())
    if word_count < 300:
        warning_msg = f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. Please enter at least 300 words for an accurate analysis."
        return "Too Short", "N/A", f"<div style='color: #b80d0d; padding: 20px; border: 1px solid #b80d0d; border-radius: 8px;'>{warning_msg}</div>", None

    try:
        tok, mod = get_model()
    except Exception as e:
        return "ERROR", "0%", f"Failed to load model: {str(e)}", None

    blocks = split_preserving_structure(text)
    pure_sents_indices = [i for i, b in enumerate(blocks) if b.strip() and not b.startswith("\n")]
    pure_sents = [blocks[i] for i in pure_sents_indices]
    
    if not pure_sents:
        return "—", "—", "<em>No sentences detected.</em>", None

    windows = []
    for i in range(len(pure_sents)):
        start = max(0, i - 1)
        end = min(len(pure_sents), i + 2)
        windows.append(" ".join(pure_sents[start:end]))

    inputs = tok(windows, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
    logits = mod(**inputs).logits
    probs = F.softmax(logits.float(), dim=-1)[:, 1].cpu().numpy().tolist()

    lengths = [len(s.split()) for s in pure_sents]
    total_words = sum(lengths)
    weighted_avg = sum(p * l for p, l in zip(probs, lengths)) / total_words if total_words > 0 else 0

    # -----------------------------
    # HTML RECONSTRUCTION
    # -----------------------------
    highlighted_html = "<div style='font-family: sans-serif; line-height: 1.8;'>"
    prob_map = {idx: probs[i] for i, idx in enumerate(pure_sents_indices)}
    
    for i, block in enumerate(blocks):
        if block.startswith("\n") or block.isspace():
            highlighted_html += block.replace("\n", "<br>")
            continue
        
        if i in prob_map:
            score = prob_map[i]
            
            # Logic: Red for > 0.81, Green for everything else (<= 0.81)
            if score >= THRESHOLD:
                color, bg = "#b80d0d", "rgba(184, 13, 13, 0.15)" # RED
            else:
                color, bg = "#11823b", "rgba(17, 130, 59, 0.15)" # GREEN
            
            highlighted_html += (
                f"<span style='background:{bg}; padding:2px 4px; border-radius:4px; border-bottom: 2px solid {color};' "
                f"title='AI Probability: {score:.1%}'>"
                f"<b style='color:{color}; font-size: 0.8em;'>[{score:.0%}]</b> {block}</span>"
            )
        else:
            highlighted_html += block
    highlighted_html += "</div>"

    # --- FINAL VERDICT ---
    if weighted_avg >= THRESHOLD:
        label = f"{weighted_avg:.0%} AI Content Detected"
        display_score = f"{weighted_avg:.1%}"
    else:
        label = "0 or * AI Content Detected"
        display_score = "*"
    
    df = pd.DataFrame({"Sentence": pure_sents, "AI Confidence": [f"{p:.1%}" for p in probs]})
    return label, display_score, highlighted_html, df

# -----------------------------
# GRADIO INTERFACE
# -----------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🕵️ AI Detector Pro")
    gr.Markdown(f"Strict Analysis. Threshold: **{THRESHOLD*100:.0f}%**. Everything below this is considered Human.")
    
    with gr.Row():
        with gr.Column(scale=3):
            text_input = gr.Textbox(label="Paste Text", lines=12, placeholder="Minimum 300 words...")
            run_btn = gr.Button("Analyze", variant="primary")
        with gr.Column(scale=1):
            verdict_out = gr.Label(label="Verdict")
            score_out = gr.Label(label="Weighted AI Score")
    
    with gr.Tabs():
        with gr.TabItem("Visual Heatmap"):
            html_out = gr.HTML()
        with gr.TabItem("Raw Data Breakdown"):
            table_out = gr.Dataframe(headers=["Sentence", "AI Confidence"], wrap=True)
            
    run_btn.click(analyze, inputs=text_input, outputs=[verdict_out, score_out, html_out, table_out])

if __name__ == "__main__":
    demo.launch()