"""
DIPPER Paraphrase API — HuggingFace Spaces deployment (GPU-enabled).
Uses google/t5-efficient-large-nl32 tokenizer with SamSJackson/paraphrase-dipper-no-ctx model.
"""
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import time

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[DIPPER] Device: {device}")

print("[DIPPER] Loading model...")
t0 = time.time()

# CRITICAL: tokenizer must come from the base model, NOT from SamSJackson repo (which has broken vocab=104)
tokenizer = AutoTokenizer.from_pretrained("google/t5-efficient-large-nl32")
model = AutoModelForSeq2SeqLM.from_pretrained("SamSJackson/paraphrase-dipper-no-ctx")
model = model.to(device)
model.eval()
print(f"[DIPPER] Model loaded in {time.time()-t0:.1f}s on {device}")
print(f"[DIPPER] Tokenizer vocab: {tokenizer.vocab_size}, Model params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")

# Self-test
test_input = "lexical = 40, order = 40 This is a test sentence about machine learning and artificial intelligence."
test_ids = tokenizer(test_input, return_tensors="pt", max_length=512, truncation=True)
test_ids = {k: v.to(device) for k, v in test_ids.items()}
with torch.no_grad():
    test_out = model.generate(**test_ids, max_new_tokens=128, top_p=0.75, do_sample=True)
test_result = tokenizer.decode(test_out[0], skip_special_tokens=True)
print(f"[DIPPER] Self-test result: '{test_result}'")

def paraphrase(text: str, lex: int = 40, order: int = 40) -> str:
    """Paraphrase a single paragraph using DIPPER."""
    if not text or not text.strip():
        return ""
    
    text = text.strip()
    prefix = f"lexical = {lex}, order = {order} "
    input_text = prefix + text
    
    inputs = tokenizer(input_text, return_tensors="pt", max_length=1000, truncation=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    t0 = time.time()
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=500,
            top_p=0.75,
            do_sample=True
        )
    elapsed = time.time() - t0
    
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"[DIPPER] {len(text)} -> {len(result)} chars in {elapsed:.1f}s")
    return result

def batch_paraphrase(text: str, lex: int = 40, order: int = 40) -> str:
    """Paraphrase multi-paragraph text by processing each paragraph separately."""
    if not text or not text.strip():
        return ""
    
    paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
    results = []
    
    for i, para in enumerate(paragraphs):
        print(f"[DIPPER] Paragraph {i+1}/{len(paragraphs)}")
        result = paraphrase(para, lex=lex, order=order)
        results.append(result)
    
    return '\n\n'.join(results)

# Gradio interface
with gr.Blocks(title="DIPPER Paraphrase API") as demo:
    gr.Markdown("# DIPPER Paraphrase API\nRewrite text to bypass AI detection. Use the API endpoint for programmatic access.")
    gr.Markdown(f"**Device:** {device} | **Tokenizer:** google/t5-efficient-large-nl32 | **Model:** SamSJackson/paraphrase-dipper-no-ctx (972M)")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text", lines=10, placeholder="Paste text here...")
            lex_slider = gr.Slider(0, 100, value=40, step=10, label="Lexical Control (40=optimal)")
            ord_slider = gr.Slider(0, 100, value=40, step=10, label="Order Control (40=optimal)")
            btn = gr.Button("Paraphrase", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Paraphrased Text", lines=10)
    
    btn.click(fn=batch_paraphrase, inputs=[input_text, lex_slider, ord_slider], outputs=output_text)

demo.launch()