from flask import Flask, request, jsonify, render_template, send_from_directory
from transformers import AutoModelForImageClassification, AutoImageProcessor
from huggingface_hub import InferenceClient
from PIL import Image
import torch
import os

app = Flask(__name__)

# =======================================
# 🔐 Hugging Face LLM Token + InferenceClient
# =======================================
HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")

client = InferenceClient(
    model="mistralai/Mistral-7B-Instruct-v0.1",
    token=HUGGINGFACE_TOKEN
)

# =======================================
# 🧠 Load Skin Disease Model
# =======================================
print("Loading skin condition classifier...")
model_name = "Jayanth2002/dinov2-base-finetuned-SkinDisease"
image_model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

# Class labels
class_names = [
    'Basal Cell Carcinoma', 'Darier_s Disease', 'Epidermolysis Bullosa Pruriginosa',
    'Hailey-Hailey Disease', 'Herpes Simplex', 'Impetigo', 'Larva Migrans',
    'Leprosy Borderline', 'Leprosy Lepromatous', 'Leprosy Tuberculoid', 'Lichen Planus',
    'Lupus Erythematosus Chronicus Discoides', 'Melanoma', 'Molluscum Contagiosum',
    'Mycosis Fungoides', 'Neurofibromatosis', 'Papilomatosis Confluentes And Reticulate',
    'Pediculosis Capitis', 'Pityriasis Rosea', 'Porokeratosis Actinic', 'Psoriasis',
    'Tinea Corporis', 'Tinea Nigra', 'Tungiasis', 'actinic keratosis', 'dermatofibroma',
    'nevus', 'pigmented benign keratosis', 'seborrheic keratosis', 'squamous cell carcinoma',
    'vascular lesion'
]

# =======================================
# 🌐 Frontend Routes
# =======================================
@app.route("/")
def index():
    return render_template("index.html")

@app.route("/upload")
def upload():
    return render_template("upload.html")

@app.route("/result")
def result():
    return render_template("result.html")  # Notice: matches the filename "results.html" instead of "result.html"

# =======================================
# 📸 /analyze Route
# =======================================
@app.route('/analyze', methods=['POST'])
def analyze():
    if 'image' not in request.files:
        return jsonify({"error": "No image uploaded"}), 400

    image_file = request.files['image']
    image = Image.open(image_file.stream).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        logits = image_model(**inputs).logits
        probs = torch.softmax(logits, dim=-1)[0]

    top_idx = torch.argmax(probs).item()
    top_conf = probs[top_idx].item()
    prediction = class_names[top_idx]

    top_conditions = sorted(
        zip(class_names, probs.tolist()),
        key=lambda x: x[1],
        reverse=True
    )[:5]

    return jsonify({
        "prediction": prediction,
        "confidence": round(top_conf, 4),
        "topConditions": [(name, round(prob, 4)) for name, prob in top_conditions],
        "description": f"{prediction} is a skin condition. Please consult a medical professional.",
        "recommendations": [
            "Take a clearer image if unsure.",
            "Consider visiting a dermatologist.",
            "Avoid self-diagnosis or self-treatment."
        ]
    })

# =======================================
# 💬 /ask Route
# =======================================
@app.route('/ask', methods=['POST'])
def ask():
    data = request.json
    question = data.get("question", "")
    condition = data.get("condition", "")

    if not question:
        return jsonify({"answer": "Please ask a valid question."}), 400

    messages = [
        {
            "role": "user",
            "content": f"A user may have {condition}. They asked: '{question}'. Respond like a helpful AI medical assistant."
        }
    ]

    try:
        response = client.chat_completion(
            messages=messages,
            max_tokens=200
        )
        answer = response.choices[0]["message"]["content"]
        return jsonify({"answer": answer.strip()})
    except Exception as e:
        return jsonify({"answer": f"Error communicating with Hugging Face: {e}"}), 500

# Add route for placeholder images if needed
@app.route('/api/placeholder/<width>/<height>')  
def placeholder(width, height):
    # This is a simple implementation - you might want to generate an actual placeholder image
    # For now, we'll just serve a static placeholder
    return send_from_directory('static/images', 'placeholder.jpg')


if __name__ == '__main__':
    app.run(debug=True)