import gradio as gr
from medrax.agent import Agent
from medrax.tools import ChestXRayClassifierTool, ChestXRaySegmentationTool, XRayVQATool
from medrax.utils import load_prompts_from_file
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver

# MedRAX ajanını başlatma fonksiyonu (main.py'den uyarlandı)
def initialize_agent(model="chatgpt-4o-latest", temperature=0.2):
    # Sistem promptunu yükle
    prompts = load_prompts_from_file("medrax/docs/system_prompts.txt")
    prompt = prompts["MEDICAL_ASSISTANT"]

    # Kullanılacak araçlar
    tools_dict = {
        "ChestXRayClassifierTool": ChestXRayClassifierTool(device="cuda"),
        "ChestXRaySegmentationTool": ChestXRaySegmentationTool(device="cuda"),
        "XRayVQATool": XRayVQATool(cache_dir="/model-weights", device="cuda"),
    }

    # Bellek ve model ayarları
    checkpointer = MemorySaver()
    model = ChatOpenAI(model=model, temperature=temperature)
    
    # Ajanı başlat
    agent = Agent(
        model,
        tools=list(tools_dict.values()),
        log_tools=True,
        log_dir="logs",
        system_prompt=prompt,
        checkpointer=checkpointer,
    )
    return agent

# Gradio arayüzü için analiz fonksiyonu
def analyze_xray(image, question):
    # Ajanı başlat
    agent = initialize_agent()
    
    # Görüntüyü ve soruyu ajana ilet
    # Not: Bu kısım MedRAX'in gerçek işleyişine bağlı olarak özelleştirilmeli
    response = agent.run(f"Analyze this chest X-ray image and answer: {question}", image=image)
    return response

# Gradio arayüzü
with gr.Blocks(title="MedRAX - Chest X-ray Analysis") as demo:
    gr.Markdown("# MedRAX: Medical Reasoning Agent for Chest X-ray")
    gr.Markdown("Upload a chest X-ray image and ask a question about it.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Chest X-ray")
            question_input = gr.Textbox(label="Your Question", placeholder="E.g., Is there a sign of pneumonia?")
            submit_btn = gr.Button("Analyze")
        with gr.Column():
            output_text = gr.Textbox(label="Analysis Result", interactive=False)
    
    # Butona tıklayınca analiz yap
    submit_btn.click(
        fn=analyze_xray,
        inputs=[image_input, question_input],
        outputs=output_text
    )

# Uygulamayı başlat
demo.launch()