import gradio as gr
from transformers import pipeline
from datetime import datetime, timezone

# Initialize the model
generator = pipeline('text-generation', model='facebook/opt-350m')  # We'll update this with your medical model later

def get_timestamp():
    """Get current UTC datetime in specified format"""
    return datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')

def format_system_info():
    """Format system information header"""
    return (
        f"Current Date and Time (UTC - YYYY-MM-DD HH:MM:SS formatted): {get_timestamp()}\n"
        f"Current User's Login: Raj-VedAI\n"
    )

def chat(message, history):
    if history is None:
        history = []
    
    # Add system information
    system_info = format_system_info()
    
    # Format the prompt with system info
    prompt = f"{system_info}\nPatient Query: {message}\nMedical AI Assistant:"

    try:
        # Generate response
        response = generator(
            prompt,
            max_length=512,
            temperature=0.7,
            top_p=0.95,
            do_sample=True
        )[0]['generated_text']
        
        # Extract and format the response
        ai_response = response[len(prompt):].strip()
        formatted_response = f"{system_info}\n{ai_response}"
        
        history.append((message, formatted_response))
        return history
    except Exception as e:
        return [(message, f"Error: {str(e)}")]

# Create custom theme
theme = gr.themes.Default().set(
    body_background_fill="#f0f8ff",  # Light blue background
    block_background_fill="#ffffff",
    block_border_width="1px",
    block_border_color="#2c3e50",
    block_radius="10px"
)

# Create the Gradio interface
demo = gr.ChatInterface(
    fn=chat,
    title="Medical Decision Support AI",
    description="""A medical decision support system that provides healthcare-related information and guidance.
    Current UTC Time: """ + get_timestamp(),
    theme=theme,
    examples=[
        "What are the symptoms of hypertension?",
        "What are common drug interactions with aspirin?",
        "What are the warning signs of diabetes?",
    ],
    retry_on_error=True
)

demo.launch()