main.py

import os
import uuid
import json
import requests
from flask import Flask, request, jsonify, send_from_directory
from werkzeug.utils import secure_filename
import google.generativeai as genai
from dotenv import load_dotenv
from flask_cors import CORS


# Step 1: API Key aur Environment Setup
load_dotenv()
try:
    genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
except TypeError:
    print("ERROR: Google API Key nahi mila.")
    print("Ek .env file banayein aur usmein 'GOOGLE_API_KEY=your_key_here' likhein.")
    exit()

app = Flask(__name__, static_folder='.', static_url_path='')
CORS(app)

# Serve static files (CSS, JS, images)
@app.route('/<path:path>')
def serve_static(path):
    return send_from_directory('.', path)

# Serve the main page
@app.route('/')
def serve_index():
    return send_from_directory('.', 'index.html')

# Step 2: API Endpoints aur Session Storage
API_ENDPOINTS = {
    "skin_disease": "https://your-api-domain.com/skin-disease-detection",
    "medicine_info": "http://localhost:5002/api/query",
    "report_reading": "https://your-api-domain.com/report-reading",
    "disease_query": "http://localhost:5001/ask"
}
SESSIONS = {}

# Step 3: Gemini se Query Classify karne ka Function
def classify_query_with_gemini(query: str):
    """User ki query ko Gemini API ka istemal karke classify karta hai."""
    model = genai.GenerativeModel('gemini-2.5-flash-lite')
    
    # *** PROMPT HAS BEEN IMPROVED ***
    prompt = f"""
    Analyze the user's medical query and classify it into one of the following categories:
    - skin_disease: For queries about skin conditions, rashes, moles, spots, or any visible symptoms on the skin.
    - medicine_info: For query about medicine(like how to use it, side effects, etc.) and also questions about a specific Medicine shown in attached image (optional) .
    - report_reading: For queries asking to interpret or explain a medical report, lab test, or blood work from an image.
    - disease_query: For general questions about diseases, symptoms, causes, or treatments.

    Based on the classification, determine if an image is essential to answer the query accurately.
    Generate response in English only.

    The user query is: 
    ---START OF QUERY---
    {query}
    ---END OF QUERY---

    Provide the output ONLY in a valid JSON format with two keys: "category" (string) and "image_required" (boolean).

    Example 1:
    Query: "what are the symptoms of typhoid"
    Output: {{"category": "disease_query", "image_required": false}}

    Example 2:
    Query: "I have a red circular rash on my arm, what is it?"
    Output: {{"category": "skin_disease", "image_required": true}}

    Example 3:
    Query: "Can you tell me what this lab report says?"
    Output: {{"category": "report_reading", "image_required": true}}

    Example 4:
    Query: "What is this white pill with 'IP 204' written on it?"
    Output: {{"category": "medicine_info", "image_required": true}}
    """
    
    try:
        response = model.generate_content(prompt)
        cleaned_text = response.text.strip().replace('```json', '').replace('```', '')
        result = json.loads(cleaned_text)
        return result
    except Exception as e:
        print(f"Gemini API call ya JSON parsing mein error: {e}")
        return None

# Step 4: API Routes (Endpoints)
@app.route('/start_session', methods=['POST'])
def start_session():
    session_id = str(uuid.uuid4())
    SESSIONS[session_id] = {"status": "started"}
    print(f"Session started: {session_id}")
    return jsonify({"session_id": session_id}), 200

@app.route('/process_query', methods=['POST'])
def process_query():
    data = request.get_json()
    session_id = data.get('session_id')
    query = data.get('query')

    if not session_id or session_id not in SESSIONS:
        return jsonify({"error": "Invalid or missing session_id"}), 400
    if not query:
        return jsonify({"error": "Query is required"}), 400

    print(f"Session {session_id}: Query received: '{query}'")
    classification = classify_query_with_gemini(query)
    
    if not classification:
        return jsonify({"error": "Could not classify the query."}), 500

    SESSIONS[session_id]['classification'] = classification
    SESSIONS[session_id]['query'] = query
    
    if classification.get('image_required'):
        print(f"Session {session_id}: Image required for category '{classification.get('category')}'")
        return jsonify({
            "status": "image_required",
            "message": "Please send the request to /process_with_image with the required photo."
        }), 200
    else:
        print(f"Session {session_id}: No image required. Forwarding to '{classification.get('category')}' API.")
        # Asli API ko call karein (Abhi ke liye mock response)
        # Session se query aur classification nikalein
        query = SESSIONS[session_id].get('query')
        classification = SESSIONS[session_id].get('classification')
        category = classification['category']
        endpoint_url = API_ENDPOINTS.get(category)
        if category == "medicine_info":
            response = requests.post(endpoint_url, json={"main_query": query})
        else:
            response = requests.post(endpoint_url, data={"main_query": query})
        del SESSIONS[session_id]
        print(f"Session {session_id} closed.")
        return jsonify({
            "status": "success",
            "response": response.json(),
            "data": f"Information about '{query}': This is a tuned response from the {classification.get('category')} service."
        })

@app.route('/process_with_image', methods=['POST'])
def process_with_image():
    session_id = request.form.get('session_id')
    
    if not session_id or session_id not in SESSIONS:
        return jsonify({"error": "Invalid or missing session_id"}), 400
    
    if 'photo' not in request.files:
        return jsonify({"error": "No photo file found in the request"}), 400
        
    file = request.files['photo']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # Session se query aur classification nikalein
    query = SESSIONS[session_id].get('query')
    classification = SESSIONS[session_id].get('classification')
    category = classification['category']
    endpoint_url = API_ENDPOINTS.get(category)
    
    print(f"Session {session_id}: Image received. Preparing to forward to '{category}' API.")

    # *** NEW: FORWARDING LOGIC THAT MATCHES YOUR CURL COMMAND ***
    # The file object from Flask needs its stream to be readable by `requests`
    # We pass the file stream, filename, and mimetype to requests
    # The dictionary key 'file' matches the '-F file=@...' part of your curl command
    files_payload = {'file': (file.filename, file.stream, file.mimetype)}
    
    # The dictionary key 'query' matches the '-F query=...' part
    data_payload = {'query': query}

    try:
        # NOTE: Neeche di gayi line asli API call hai.
        # Jab aapka backend service (e.g., http://localhost:5002/api/query) taiyaar ho,
        # to is line ko uncomment kar dein.
        
        response_from_service = requests.post(endpoint_url, files=files_payload, data=data_payload)
        response_from_service.raise_for_status() # Agar 4xx/5xx error ho to exception raise karega
        tuned_response = response_from_service.json() # Assume service returns JSON

        # Abhi ke liye, hum ek mock response bhej rahe hain
        mock_response = {
            "status": "success",
            "response": tuned_response,
            "data": f"Analysis for '{query}' based on your image: This is a tuned MOCK response from the {category} service."
        }
        
        # Session close karein
        del SESSIONS[session_id]
        print(f"Session {session_id} closed.")
        
        return jsonify(mock_response)

    except requests.exceptions.RequestException as e:
        del SESSIONS[session_id]
        print(f"Session {session_id} closed after failed API call.")
        return jsonify({"status": "error", "message": f"Backend service call failed: {e}"}), 503