import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from datetime import datetime, timezone
import os
from huggingface_hub import login
from huggingface_hub.utils import RepositoryNotFoundError, HfHubHTTPError

def initialize_model():
    try:
        # Login to Hugging Face Hub
        token = os.getenv("HUGGING_FACE_HUB_TOKEN")
        if not token:
            return False, "No token found. Please set HUGGING_FACE_HUB_TOKEN in Space secrets.", None
            
        login(token=token)
        
        # Initialize the model and tokenizer with token (not use_auth_token)
        model_id = "CohereLabs/c4ai-command-a-03-2025"
        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            token=token  # Updated from use_auth_token
        )
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            token=token  # Updated from use_auth_token
        )
        return True, model, tokenizer
    except RepositoryNotFoundError:
        return False, "Model repository not found. Please check the model ID.", None
    except HfHubHTTPError as e:
        if e.response.status_code == 401:
            return False, "Authentication failed. Please check your token permissions.", None
        elif e.response.status_code == 403:
            return False, "Access denied. Please request access at https://huggingface.co/CohereLabs/c4ai-command-a-03-2025", None
        else:
            return False, f"An error occurred: {str(e)}", None
    except Exception as e:
        return False, f"Unexpected error: {str(e)}", None

# Initialize model and handle potential errors
success, result, tokenizer = initialize_model()
if not success:
    print(f"Error initializing model: {result}")
else:
    model = result

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

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 not success:
        return [(message, f"Error: {result}")]
    
    if history is None:
        history = []
    
    try:
        # Add system information
        system_info = format_system_info()
        
        # Format messages with the c4ai-command-a-03-2025 chat template
        messages = [{"role": "user", "content": message}]
        input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
        
        # Generate response
        gen_tokens = model.generate(
            input_ids,
            max_new_tokens=100,
            do_sample=True,
            temperature=0.3,
        )
        
        # Decode response
        gen_text = tokenizer.decode(gen_tokens[0])
        
        # Format the full response with system info
        formatted_response = f"{system_info}\n{gen_text}"
        history.append((message, formatted_response))
        return history
    except Exception as e:
        return [(message, f"Error during chat: {str(e)}")]

# Create the Gradio interface - removed retry_on_error
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(),
    examples=[
        "What are the symptoms of hypertension?",
        "What are common drug interactions with aspirin?",
        "What are the warning signs of diabetes?",
    ]
)

demo.launch()