import gradio as gr
from transformers import pipeline

# Load Hugging Face models for text generation, sentiment analysis, and zero-shot classification
# These models handle the core AI processes from the BRD: message generation, response evaluation for risk, sentiment, and severity
text_generator = pipeline("text-generation", model="gpt2")  # For generating follow-up messages dynamically (replaces FR1's automation logic, but focuses on generation instead of sending)
sentiment_analyzer = pipeline("sentiment-analysis")  # For evaluating sentiment (FR8)
zero_shot_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")  # For risk classification and severity (FR2, FR3, FR8)

def generate_followup_message(prompt):
    """
    Generates a follow-up message based on a prompt using Hugging Face text-generation model.
    This implements the message creation part from the BRD (adapted from FR1, FR6, FR7 - supports dynamic, multilingual-capable generation).
    """
    generated = text_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
    return generated

def analyze_patient_response(response):
    """
    Analyzes patient response for risk, severity, and sentiment using Hugging Face models.
    This covers BRD processes: FR2 (store/evaluate for risk), FR3 (trigger based on high-risk), FR8 (sentiment and severity).
    Risk triage: Uses zero-shot classification with candidate labels for low/medium/high risk.
    Sentiment: Positive/Negative with score.
    Severity: Derived from top risk label.
    Note: In a full system, high-risk would trigger alerts; here it's analyzed for demo.
    """
    # Sentiment analysis
    sentiment_result = sentiment_analyzer(response)[0]
    sentiment = sentiment_result['label']
    sentiment_score = sentiment_result['score']

    # Risk classification (zero-shot for flexibility, can be fine-tuned in production)
    candidate_labels = ["low risk symptoms (e.g., mild discomfort)", "medium risk symptoms (e.g., moderate pain)", "high risk symptoms (e.g., severe symptoms requiring immediate attention)"]
    risk_result = zero_shot_classifier(response, candidate_labels=candidate_labels)
    risk_level = risk_result['labels'][0]  # Highest scoring label
    risk_score = max(risk_result['scores'])  # Score of the top label
    severity = "High" if "high" in risk_level else "Medium" if "medium" in risk_level else "Low"

    # Simulate trigger for high-risk (FR3): In full app, this could create a 'case' or alert
    risk_alert = "High-risk detected: Escalate to case creation." if "high" in risk_level else "No immediate escalation needed."

    return {
        "sentiment": sentiment,
        "sentiment_score": sentiment_score,
        "risk_level": risk_level,
        "risk_score": risk_score,
        "severity": severity,
        "risk_alert": risk_alert
    }

# Gradio UI for the Hugging Face Space app
# This provides a dashboard-like interface for clinic staff (FR5), with tabs for generation and analysis
# Replaces Salesforce dashboards/LWC with a simple web UI
# Note: Consent verification (FR4) and multilingual support (FR7) can be added via prompts/models; scalability/performance handled by HF hosting
with gr.Blocks(title="AI-Powered Patient Follow-up Agent (Hugging Face Implementation)") as demo:
    gr.Markdown("""
    # AI-Powered Patient Follow-up Agent
    This app implements the core processes from the BRD using Hugging Face models:
    - Generate dynamic follow-up messages (adapted from FR1, without actual sending).
    - Analyze responses for risk, sentiment, and severity (FR2, FR3, FR8).
    - Focus on AI logic; assumes integration with external tools for sending/logging if needed.
    - HIPAA compliance: Ensure secure deployment; this is a prototype.
    """)

    with gr.Tab("Generate Follow-up Message"):
        prompt_input = gr.Textbox(label="Input Prompt (e.g., 'Generate a follow-up message for a patient after surgery in Spanish')", lines=3)
        generate_button = gr.Button("Generate Message")
        message_output = gr.Textbox(label="Generated Follow-up Message", lines=5)
        generate_button.click(fn=generate_followup_message, inputs=prompt_input, outputs=message_output)

    with gr.Tab("Analyze Patient Response"):
        response_input = gr.Textbox(label="Patient Response Text", lines=3)
        analyze_button = gr.Button("Analyze Response")
        analysis_output = gr.JSON(label="Analysis Results (Risk, Sentiment, Severity)")
        analyze_button.click(fn=analyze_patient_response, inputs=response_input, outputs=analysis_output)

demo.launch()