import streamlit as st
from PIL import Image
from huggingface_hub import InferenceClient
import io
import base64

# --- Configuration (Simplified for Spaces) ---

# No need for API token if running *within* a Space
# The Space's environment will handle authentication

# --- Image Encoding ---
def encode_image(image):
    buffered = io.BytesIO()
    # Convert to RGB *before* saving as JPEG
    if image.mode == "RGBA":
        image = image.convert("RGB")
    image.save(buffered, format="JPEG")  # Save as JPEG
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str


# --- Model Interaction (using InferenceClient) ---

def analyze_image_with_maira(image):
    """Analyzes the image using the Maira-2 model via the Hugging Face Inference API.
    """
    try:
        encoded_image = encode_image(image)
        client = InferenceClient()  # No token needed inside the Space
        result = client.question_answering(
            question="Analyze this chest X-ray image and provide detailed findings. Include any abnormalities, their locations, and potential diagnoses. Be as specific as possible.",
            image=encoded_image,  # Pass the encoded image directly
            model="microsoft/maira-2"  # Specify the model
        )
        return result

    except Exception as e:
        st.error(f"An error occurred: {e}")  # General exception handling is sufficient here
        return None


# --- Streamlit App ---

def main():
    st.title("Chest X-ray Analysis with Maira-2 (Hugging Face Spaces)")
    st.write(
        "Upload a chest X-ray image. This app uses the Maira-2 model within this Hugging Face Space."
    )

    uploaded_file = st.file_uploader("Choose a chest X-ray image (JPG, PNG)", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        with st.spinner("Analyzing image with Maira-2..."):
            analysis_results = analyze_image_with_maira(image)

        if analysis_results:
            # --- Results Display (VQA format) ---
            if isinstance(analysis_results, dict) and 'answer' in analysis_results:
                st.subheader("Findings:")
                st.write(analysis_results['answer'])
            else:
                st.warning("Unexpected API response format.")
                st.write("Raw API response:", analysis_results)
        else:
            st.error("Failed to get analysis results.")

    else:
        st.write("Please upload an image.")

    st.write("---")
    st.write("Disclaimer: For informational purposes only.  Not medical advice.")

if __name__ == "__main__":
    main()