import gradio as gr
from huggingface_hub import InferenceApi
from transformers import pipeline
import requests
from PIL import Image
import io
import os
import openai

# -----------------------------
# CONFIGURATION
# -----------------------------
# Hugging Face model for acne classification
MODEL_ID = "imfarzanansari/skintelligent-acne"

# Use local pipeline for image classification (faster and more stable)
classifier = pipeline("image-classification", model=MODEL_ID)

# Set your Mistral API key (via environment variable)
openai.api_key = os.getenv("MISTRAL_API_KEY")

# -----------------------------
# HELPER FUNCTIONS
# -----------------------------
def classify_acne(image_url):
    try:
        response = requests.get(image_url)
        img = Image.open(io.BytesIO(response.content)).convert("RGB")
    except Exception as e:
        return "❌ Could not load image. Please check the URL.", "", None

    # Run the acne classification
    preds = classifier(img)
    if not preds:
        return "No prediction.", "", img

    top_pred = preds[0]["label"]
    score = preds[0]["score"]

    # Explanation text
    explanation = explain_acne_type(top_pred)

    result_text = f"**Detected Acne Type:** {top_pred}\n\n**Confidence:** {score:.2f}"
    return result_text, explanation, img


def explain_acne_type(acne_type):
    explanations = {
        "Blackheads": "Blackheads are open comedones caused by clogged hair follicles with sebum and dead skin. They appear black due to oxidation.",
        "Whiteheads": "Whiteheads are closed comedones formed when pores are clogged with oil and dead skin but remain closed at the surface.",
        "Papules": "Papules are small, red, inflamed bumps without visible pus. They often result from irritated or infected clogged pores.",
        "Pustules": "Pustules are pus-filled pimples with a white or yellow center. They can be tender and are often caused by bacterial infection.",
        "Nodules": "Nodules are large, painful lumps deep under the skin caused by severe inflammation and infection in clogged pores.",
        "Cysts": "Cysts are severe acne lesions filled with pus and can cause scarring if not treated properly.",
    }
    return explanations.get(acne_type, "This acne type is uncommon or not specifically defined in the dataset.")


def query_acne_info(acne_type, user_query):
    if not user_query.strip():
        return "Please enter a question."
    try:
        prompt = f"You are an expert dermatologist. The user has acne type '{acne_type}'. Answer this query:\n{user_query}"
        completion = openai.ChatCompletion.create(
            model="mistral-tiny",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.6,
        )
        return completion.choices[0].message["content"]
    except Exception as e:
        return f"Error: {str(e)}"


# -----------------------------
# GRADIO INTERFACE
# -----------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="Acne Type Classifier & Chatbot") as demo:
    gr.Markdown(
        """
        # 🧴 Acne Type Classifier & Dermatology Assistant  
        Upload or paste the URL of an acne image to detect its type.  
        Then ask any query about the detected acne type using the chatbot below.
        """
    )

    with gr.Row():
        image_url = gr.Textbox(label="🔗 Enter Image URL", placeholder="Paste image URL here...")
        submit_btn = gr.Button("Classify", variant="primary")

    with gr.Row():
        with gr.Column(scale=1):
            image_output = gr.Image(label="Uploaded Image", type="pil")

        with gr.Column(scale=2):
            result_box = gr.Markdown(label="Prediction Result", elem_classes="big-box")
            explanation_box = gr.Textbox(
                label="Acne Explanation",
                lines=6,
                interactive=False,
                elem_classes="big-box"
            )

    # Chatbot section
    with gr.Accordion("💬 Ask Dermatology Chatbot", open=True):
        with gr.Row():
            user_query = gr.Textbox(
                label="Enter your query about the detected acne",
                placeholder="e.g., What is the best treatment for cystic acne?",
                lines=2,
            )
        with gr.Row():
            chat_response = gr.Textbox(
                label="Chatbot Response",
                lines=6,
                interactive=False,
                elem_classes="big-box"
            )
        chat_btn = gr.Button("Ask Chatbot", variant="secondary")

    # Button functionality
    submit_btn.click(
        classify_acne,
        inputs=[image_url],
        outputs=[result_box, explanation_box, image_output],
    )

    chat_btn.click(
        query_acne_info,
        inputs=[result_box, user_query],
        outputs=[chat_response],
    )

    gr.Markdown(
        "#### ⚕️ Disclaimer: This app provides general information and should not replace professional medical advice."
    )

# Custom CSS to enlarge boxes
demo.css = """
.big-box textarea, .big-box pre, .big-box .wrap {
    height: 220px !important;
    font-size: 16px;
}
"""

# Launch app
if __name__ == "__main__":
    demo.launch()