import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import warnings

warnings.filterwarnings("ignore")

model = None
tokenizer = None
model_loaded = False

def load_model():
    global model, tokenizer, model_loaded
    
    if model_loaded:
        return True
    
    try:
        print("[INFO] Model loading...")
        
        model_path = "Neurazum/Lbai-1-preview"
        
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float32,
            device_map="cpu",
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        
        model.eval()
        model_loaded = True
        print("[INFO] Model successfully loaded!")
        return True
        
    except Exception as e:
        print(f"[ERROR] Model yüklenemedi: {e}")
        return False

def respond(message, history, system_message, max_tokens, temperature, top_p):
    global model, tokenizer, model_loaded
    
    if not message or message.strip() == "":
        yield "Please write a message..."
        return
    
    if not model_loaded:
        yield "⏳ Loading model, please wait..."
        if not load_model():
            yield "❌ Model could not be loaded. Please try again later."
            return
    
    try:
        prompt = f"{system_message}\n\n"
        
        if history:
            for item in history:
                try:
                    if isinstance(item, (list, tuple)) and len(item) >= 2:
                        user_msg, assistant_msg = item[0], item[1]
                        if user_msg:
                            prompt += f"Patient: {user_msg}\n"
                        if assistant_msg:
                            prompt += f"Doctor: {assistant_msg}\n"
                    elif isinstance(item, dict):
                        role = item.get("role", "")
                        content = item.get("content", "")
                        if role == "user" and content:
                            prompt += f"Patient: {content}\n"
                        elif role == "assistant" and content:
                            prompt += f"Doctor: {content}\n"
                except Exception as e:
                    print(f"[WARNING] History item skipped: {e}")
                    continue
        
        prompt += f"Patient: {message}\nDoctor:"
        
        print(f"[DEBUG] Prompt: {prompt[:300]}...")
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        if "Doctor:" in full_response:
            response = full_response.split("Doctor:")[-1].strip()
        else:
            response = full_response[len(prompt):].strip()
        
        if "\nPatient:" in response:
            response = response.split("\nPatient:")[0].strip()
        
        yield response if response else "The model could not generate a response."
        
    except Exception as e:
        import traceback
        print(f"[ERROR] {traceback.format_exc()}")
        yield f"❌ Hata: {str(e)}"

chatbot = gr.ChatInterface(
    fn=respond,
    title="Lbai-1-preview",
    description="Ask your medical questions. The model will load in the first message, so please wait a moment. Artificial intelligence can make mistakes.",
    additional_inputs=[
        gr.Textbox(
            value="You are a helpful medical assistant.", 
            label="System Prompt",
            lines=2
        ),
        gr.Slider(
            minimum=1, 
            maximum=512, 
            value=200, 
            step=1, 
            label="Max Tokens"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=2.0, 
            value=0.7, 
            step=0.1, 
            label="Temperature"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top-p",
        ),
    ],
)

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