import gradio as gr
from transformers import pipeline
import torch

# Model ID for Microsoft's Phi-3 (Ungated, high-efficiency SLM)
model_id = "microsoft/Phi-3-mini-4k-instruct"

print("Initializing Digital Vaccine Pipeline...")

# Initialize the pipeline
# We use trust_remote_code=True because Phi-3 uses a custom architecture
pipe = pipeline(
    "text-generation",
    model=model_id,
    device_map="auto",
    model_kwargs={
        "torch_dtype": "auto",
        "trust_remote_code": True,
        "attn_implementation": "eager"  # <--- THIS IS THE FIX
    }
)
def digital_wellness_intervention(user_input, scroll_velocity):
    """
    Determines if the user is in a 'Zombiescroll' state and 
    generates a clinical AI intervention.
    """
    # Threshold for intervention (Simulated Telemetry)
    is_zombiescroll = scroll_velocity > 70 
    
    # Building the Phi-3 prompt template
    if is_zombiescroll:
        system_msg = "You are a clinical AI specialist focused on digital wellness."
        user_msg = f"The user is currently 'doomscrolling' with a high velocity of {scroll_velocity}. Their last stated intent was: '{user_input}'. Provide a 2-sentence Micro-CBT intervention to help them regain consciousness and stop scrolling."
    else:
        system_msg = "You are a supportive digital wellness assistant."
        user_msg = f"The user is browsing intentionally. Their intent is: '{user_input}'. Give a very brief (1 sentence) encouragement to stay focused."

    # Phi-3 specific formatting
    prompt = f"<|system|>\n{system_msg}<|end|>\n<|user|>\n{user_msg}<|end|>\n<|assistant|>\n"
    
    # Generate response
    outputs = pipe(
        prompt, 
        max_new_tokens=100, 
        do_sample=True, 
        temperature=0.7,
        return_full_text=False
    )
    
    intervention_text = outputs[0]["generated_text"].strip()
    
    status = "⚠️ INTERVENTION TRIGGERED" if is_zombiescroll else "✅ FOCUS MAINTAINED"
    return status, intervention_text

# --- Gradio UI Setup ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧬 The Digital Vaccine: AI-in-the-Loop")
    gr.Markdown(
        "### A Technical Prototype for Behavioral Intervention\n"
        "This system uses a **Small Language Model (Phi-3)** to act as a 'Circuit Breaker' for digital addiction. "
        "By monitoring simulated scroll telemetry, the AI provides real-time clinical nudges."
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### User Telemetry")
            input_text = gr.Textbox(
                label="What is your current goal?", 
                placeholder="e.g., Researching LLM quantization..."
            )
            velocity = gr.Slider(
                minimum=0, 
                maximum=100, 
                value=20, 
                label="Scroll Velocity (Dopamine Intensity)"
            )
            btn = gr.Button("Analyze Mental State", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("### AI Diagnostic")
            status_output = gr.Label(label="Current State")
            intervention_output = gr.Textbox(label="MedGemma-style Intervention", lines=4)

    btn.click(
        digital_wellness_intervention, 
        inputs=[input_text, velocity], 
        outputs=[status_output, intervention_output]
    )
    
    gr.Examples(
        examples=[
            ["I'm looking for a recipe for dinner", 25],
            ["Just killing time, feeling bored", 85],
            ["Checking work emails", 40]
        ],
        inputs=[input_text, velocity]
    )

    gr.Markdown("---")
    gr.Markdown("*Technical Note: This prototype utilizes Microsoft's Phi-3-mini (3.8B) for edge-compatible inference.*")

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