import gradio as gr
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    SynthIDTextWatermarkingConfig,
    SynthIDTextBayesianDetector
)

# Initialize model and tokenizer
MODEL_NAME = "google/gemma-2b"  # You can change this to your preferred model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Configure watermarking
WATERMARK_KEYS = [654, 400, 836, 123, 340, 443, 597, 160, 57, 789]  # Example keys
watermarking_config = SynthIDTextWatermarkingConfig(
    keys=WATERMARK_KEYS,
    ngram_len=5
)

# Initialize detector
detector = SynthIDTextBayesianDetector(watermarking_config)

def apply_watermark(text):
    """Apply SynthID watermark to input text."""
    try:
        # Tokenize input
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        
        # Generate with watermark
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                watermarking_config=watermarking_config,
                do_sample=True,
                max_length=len(inputs["input_ids"][0]) + 100,  # Add some extra tokens
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode output
        watermarked_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return watermarked_text, "Watermark applied successfully!"
    except Exception as e:
        return text, f"Error applying watermark: {str(e)}"

def detect_watermark(text):
    """Detect if text contains SynthID watermark."""
    try:
        # Get detection score
        score = detector.detect(text)
        
        # Interpret results
        threshold = 0.5  # You can adjust this threshold
        is_watermarked = score > threshold
        
        result = f"Watermark Detection Score: {score:.3f}\n"
        result += f"Verdict: {'WATERMARK DETECTED' if is_watermarked else 'NO WATERMARK DETECTED'}"
        
        return result
    except Exception as e:
        return f"Error detecting watermark: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
    gr.Markdown("# SynthID Text Watermarking Tool")
    gr.Markdown("Apply and detect SynthID watermarks in text")
    
    with gr.Tab("Apply Watermark"):
        with gr.Row():
            input_text = gr.Textbox(label="Input Text", lines=5)
            output_text = gr.Textbox(label="Watermarked Text", lines=5)
            status = gr.Textbox(label="Status")
        apply_btn = gr.Button("Apply Watermark")
        apply_btn.click(apply_watermark, inputs=[input_text], outputs=[output_text, status])
    
    with gr.Tab("Detect Watermark"):
        with gr.Row():
            detect_input = gr.Textbox(label="Text to Check", lines=5)
            detect_result = gr.Textbox(label="Detection Result", lines=3)
        detect_btn = gr.Button("Detect Watermark")
        detect_btn.click(detect_watermark, inputs=[detect_input], outputs=[detect_result])
    
    gr.Markdown("""
    ### Notes:
    - The watermark is designed to be imperceptible to humans but detectable by the classifier
    - Detection scores above 0.5 indicate likely presence of a watermark
    - The watermark is somewhat robust to minor text modifications but may not survive major changes
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()