from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from typing import Dict, List, Tuple
import uuid

MODEL_NAME = "BioMistral/BioMistral-7B"

print("🔹 Loading BioMistral model... This may take a while on first run.")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    torch_dtype="auto"
)

chat_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9
)

# Store conversation history per session
conversation_sessions: Dict[str, List[Dict[str, str]]] = {}

SYSTEM_PROMPT = (
    "You are MediScope AI, a medical assistant that helps patients understand "
    "their symptoms in simple, safe, and educational language. "
    "Always encourage professional consultation for serious conditions."
)

def get_or_create_session(session_id: str) -> List[Dict[str, str]]:
    """Get existing session or create a new one."""
    if session_id not in conversation_sessions:
        conversation_sessions[session_id] = []
    return conversation_sessions[session_id]

def build_conversation_prompt(history: List[Dict[str, str]], user_prompt: str) -> str:
    """Build the full conversation prompt from history and new user message."""
    prompt_parts = [SYSTEM_PROMPT]
    
    # Add conversation history
    for msg in history:
        if msg["role"] == "user":
            prompt_parts.append(f"User: {msg['content']}")
        elif msg["role"] == "assistant":
            prompt_parts.append(f"Assistant: {msg['content']}")
    
    # Add current user message
    prompt_parts.append(f"User: {user_prompt}")
    prompt_parts.append("Assistant:")
    
    return "\n\n".join(prompt_parts)

def chat_with_biomistral(user_prompt: str, session_id: str = None) -> Tuple[str, str]:
    """
    Chat with BioMistral model, maintaining conversation history.
    
    Args:
        user_prompt: The user's message
        session_id: Optional session ID. If None, a new session is created.
    
    Returns:
        tuple: (response_text, session_id)
    """
    # Generate or use provided session ID
    if session_id is None:
        session_id = str(uuid.uuid4())
    
    # Get conversation history for this session
    history = get_or_create_session(session_id)
    
    # Build the full conversation prompt
    full_prompt = build_conversation_prompt(history, user_prompt)
    
    # Generate response
    response = chat_pipeline(full_prompt)[0]["generated_text"]
    
    # Extract only the assistant's reply (everything after the last "Assistant:")
    reply = response.split("Assistant:")[-1].strip()
    
    # Update conversation history
    history.append({"role": "user", "content": user_prompt})
    history.append({"role": "assistant", "content": reply})
    
    return reply, session_id

def clear_session(session_id: str) -> bool:
    """Clear conversation history for a session."""
    if session_id in conversation_sessions:
        conversation_sessions[session_id] = []
        return True
    return False