import gradio as gr
import torch
import requests
import tempfile
from pathlib import Path
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor

_model_cache = {}

def load_model_and_processor(hf_token: str):
    """
    Loads the MAIRA-2 model and processor from Hugging Face using the provided token.
    The loaded objects are cached keyed by the token.
    """
    if hf_token in _model_cache:
        return _model_cache[hf_token]
    device = torch.device("cpu")
    model = AutoModelForCausalLM.from_pretrained(
        "microsoft/maira-2",
        trust_remote_code=True,
        use_auth_token=hf_token
    )
    processor = AutoProcessor.from_pretrained(
        "microsoft/maira-2",
        trust_remote_code=True,
        use_auth_token=hf_token
    )
    model.eval()
    model.to(device)
    _model_cache[hf_token] = (model, processor)
    return model, processor

def get_sample_data() -> dict:
    """
    Downloads sample chest X-ray images and associated data.
    """
    frontal_image_url = "https://openi.nlm.nih.gov/imgs/512/145/145/CXR145_IM-0290-1001.png"
    lateral_image_url = "https://openi.nlm.nih.gov/imgs/512/145/145/CXR145_IM-0290-2001.png"

    def download_and_open(url: str) -> Image.Image:
        response = requests.get(url, headers={"User-Agent": "MAIRA-2"}, stream=True)
        return Image.open(response.raw).convert("RGB")

    frontal = download_and_open(frontal_image_url)
    lateral = download_and_open(lateral_image_url)
    return {
        "frontal": frontal,
        "lateral": lateral,
        "indication": "Dyspnea.",
        "technique": "PA and lateral views of the chest.",
        "comparison": "None.",
        "phrase": "Pleural effusion."
    }

def generate_report(hf_token, frontal, lateral, indication, technique, comparison, use_grounding):
    """
    Generates a radiology report using the MAIRA-2 model.
    If any image/text input is missing, sample data is used.
    """
    try:
        model, processor = load_model_and_processor(hf_token)
    except Exception as e:
        return f"Error loading model: {str(e)}"
    device = torch.device("cpu")
    sample = get_sample_data()
    if frontal is None:
        frontal = sample["frontal"]
    if lateral is None:
        lateral = sample["lateral"]
    if not indication:
        indication = sample["indication"]
    if not technique:
        technique = sample["technique"]
    if not comparison:
        comparison = sample["comparison"]

    processed_inputs = processor.format_and_preprocess_reporting_input(
        current_frontal=frontal,
        current_lateral=lateral,
        prior_frontal=None,  # No prior study is used in this demo.
        indication=indication,
        technique=technique,
        comparison=comparison,
        prior_report=None,
        return_tensors="pt",
        get_grounding=use_grounding,
    )
    # Move all tensors to the CPU
    processed_inputs = {k: v.to(device) for k, v in processed_inputs.items()}
    # Remove keys containing "image_sizes" to prevent unexpected keyword errors.
    processed_inputs = dict(processed_inputs)
    keys_to_remove = [k for k in processed_inputs if "image_sizes" in k]
    for key in keys_to_remove:
        processed_inputs.pop(key, None)
    
    max_tokens = 450 if use_grounding else 300
    with torch.no_grad():
        output_decoding = model.generate(
            **processed_inputs,
            max_new_tokens=max_tokens,
            use_cache=True,
        )
    prompt_length = processed_inputs["input_ids"].shape[-1]
    decoded_text = processor.decode(output_decoding[0][prompt_length:], skip_special_tokens=True)
    decoded_text = decoded_text.lstrip()  # Remove any leading whitespace
    prediction = processor.convert_output_to_plaintext_or_grounded_sequence(decoded_text)
    return prediction

def run_phrase_grounding(hf_token, frontal, phrase):
    """
    Runs phrase grounding using the MAIRA-2 model.
    If image or phrase is missing, sample data is used.
    """
    try:
        model, processor = load_model_and_processor(hf_token)
    except Exception as e:
        return f"Error loading model: {str(e)}"
    device = torch.device("cpu")
    sample = get_sample_data()
    if frontal is None:
        frontal = sample["frontal"]
    if not phrase:
        phrase = sample["phrase"]
    processed_inputs = processor.format_and_preprocess_phrase_grounding_input(
        frontal_image=frontal,
        phrase=phrase,
        return_tensors="pt",
    )
    processed_inputs = {k: v.to(device) for k, v in processed_inputs.items()}
    # Remove keys containing "image_sizes" to prevent unexpected keyword errors.
    processed_inputs = dict(processed_inputs)
    keys_to_remove = [k for k in processed_inputs if "image_sizes" in k]
    for key in keys_to_remove:
        processed_inputs.pop(key, None)
    
    with torch.no_grad():
        output_decoding = model.generate(
            **processed_inputs,
            max_new_tokens=150,
            use_cache=True,
        )
    prompt_length = processed_inputs["input_ids"].shape[-1]
    decoded_text = processor.decode(output_decoding[0][prompt_length:], skip_special_tokens=True)
    prediction = processor.convert_output_to_plaintext_or_grounded_sequence(decoded_text)
    return prediction


def login_ui(hf_token):
    """Authenticate the user by loading the model."""
    try:
        load_model_and_processor(hf_token)
        return "🔓 Login successful! You can now use the model."
    except Exception as e:
        return f"❌ Login failed: {str(e)}"

def generate_report_ui(hf_token, frontal_path, lateral_path, indication, technique, comparison,
                         prior_frontal_path, prior_lateral_path, prior_report, grounding):
    """
    Wrapper for generate_report that accepts file paths (from the UI) for images.
    Prior study fields are ignored.
    """
    try:
        frontal = Image.open(frontal_path) if frontal_path else None
        lateral = Image.open(lateral_path) if lateral_path else None
    except Exception as e:
        return f"❌ Error loading images: {str(e)}"
    return generate_report(hf_token, frontal, lateral, indication, technique, comparison, grounding)

def run_phrase_grounding_ui(hf_token, frontal_path, phrase):
    """
    Wrapper for run_phrase_grounding that accepts a file path for the frontal image.
    """
    try:
        frontal = Image.open(frontal_path) if frontal_path else None
    except Exception as e:
        return f"❌ Error loading image: {str(e)}"
    return run_phrase_grounding(hf_token, frontal, phrase)

def save_temp_image(img: Image.Image) -> str:
    """Save a PIL image to a temporary file and return the file path."""
    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    img.save(temp_file.name)
    return temp_file.name

def load_sample_findings():
    """
    Loads sample data for the report generation tab.
    Returns file paths for current study images, sample text fields, and dummy values for prior study.
    """
    sample = get_sample_data()
    return [
        save_temp_image(sample["frontal"]),  # frontal image file path
        save_temp_image(sample["lateral"]),    # lateral image file path
        sample["indication"],
        sample["technique"],
        sample["comparison"],
        None,  # prior frontal (not used)
        None,  # prior lateral (not used)
        None,  # prior report (not used)
        False  # grounding checkbox default
    ]

def load_sample_phrase():
    """
    Loads sample data for the phrase grounding tab.
    Returns file path for the frontal image and a sample phrase.
    """
    sample = get_sample_data()
    return [save_temp_image(sample["frontal"]), sample["phrase"]]


with gr.Blocks(title="MAIRA-2 Medical Assistant") as demo:
    gr.Markdown(
        """
        # MAIRA-2 Medical Assistant
        **Authentication required** - You need a Hugging Face account and access token to use this model.
        1. Get your access token from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
        2. Request model access at [https://huggingface.co/microsoft/maira-2](https://huggingface.co/microsoft/maira-2)
        3. Paste your token below to begin
        """
    )
    
    with gr.Row():
        hf_token = gr.Textbox(
            label="Hugging Face Token",
            placeholder="hf_xxxxxxxxxxxxxxxxxxxx",
            type="password"
        )
        login_btn = gr.Button("Authenticate")
        login_status = gr.Textbox(label="Authentication Status", interactive=False)
    
    login_btn.click(
        login_ui,
        inputs=hf_token,
        outputs=login_status
    )
    
    with gr.Tabs():
        with gr.Tab("Report Generation"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("## Current Study")
                    frontal = gr.Image(label="Frontal View", type="filepath")
                    lateral = gr.Image(label="Lateral View", type="filepath")
                    indication = gr.Textbox(label="Clinical Indication")
                    technique = gr.Textbox(label="Imaging Technique")
                    comparison = gr.Textbox(label="Comparison")
                    
                    gr.Markdown("## Prior Study (Optional)")
                    prior_frontal = gr.Image(label="Prior Frontal View", type="filepath")
                    prior_lateral = gr.Image(label="Prior Lateral View", type="filepath")
                    prior_report = gr.Textbox(label="Prior Report")
                    
                    grounding = gr.Checkbox(label="Include Grounding")
                    sample_btn = gr.Button("Load Sample Data")
                with gr.Column():
                    report_output = gr.Textbox(label="Generated Report", lines=10)
                    generate_btn = gr.Button("Generate Report")
            
            sample_btn.click(
                load_sample_findings,
                outputs=[frontal, lateral, indication, technique, comparison,
                         prior_frontal, prior_lateral, prior_report, grounding]
            )
            generate_btn.click(
                generate_report_ui,
                inputs=[hf_token, frontal, lateral, indication, technique, comparison,
                        prior_frontal, prior_lateral, prior_report, grounding],
                outputs=report_output
            )
        
        with gr.Tab("Phrase Grounding"):
            with gr.Row():
                with gr.Column():
                    pg_frontal = gr.Image(label="Frontal View", type="filepath")
                    phrase = gr.Textbox(label="Phrase to Ground")
                    pg_sample_btn = gr.Button("Load Sample Data")
                with gr.Column():
                    pg_output = gr.Textbox(label="Grounding Result", lines=3)
                    pg_btn = gr.Button("Find Phrase")
            
            pg_sample_btn.click(
                load_sample_phrase,
                outputs=[pg_frontal, phrase]
            )
            pg_btn.click(
                run_phrase_grounding_ui,
                inputs=[hf_token, pg_frontal, phrase],
                outputs=pg_output
            )

demo.launch()