import os
import io
import base64
import time
import logging
import re
import sys
from pathlib import Path
from PIL import Image
import gradio as gr
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager

# إعداد Logging شامل
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

logger.info("🚀 بدء تشغيل DermaScan AI...")

try:
    logger.info("📦 جاري استيراد الوحدات...")
    from rag import build_default_vectorstore
    from agent import (
        build_agent,
        set_session_storage,
        set_current_thread_id,
        set_api_client,
    )
    from api_client import DermaScanAPIClient
    logger.info("✅ تم استيراد جميع الوحدات بنجاح")
except Exception as e:
    logger.error(f"❌ فشل استيراد الوحدات: {e}", exc_info=True)
    raise

# Configuration
MODEL_NAME = "llama-3.1-8b-instant"
FALLBACK_MODEL_NAME = "llama-3.3-70b-versatile"
TEMPERATURE = 0.25
TOP_K = 3
EXTERNAL_API_URL = "https://omarelrayes-api.hf.space"

# Global variables
SESSIONS = {}
agent = None
fallback_agent = None
api_client = None

def get_session(thread_id: str) -> dict:
    if thread_id not in SESSIONS:
        SESSIONS[thread_id] = {
            "role": "patient",
            "last_analysis": None,
            "current_image": None,
            "matched_patient_id": None,
            "matched_patient_report": None,
            "_image_blobs": {},
            "_session_image_names": [],
            "_turn_count": 0,
            "_summary": "",
        }
    return SESSIONS[thread_id]

def initialize_system():
    """Initialize agent and API client"""
    global agent, fallback_agent, api_client
    
    try:
        logger.info("🔧 بدء تهيئة النظام...")
        
        # Check Groq API Key
        groq_key = os.environ.get("GROQ_API_KEY")
        if not groq_key:
            logger.warning("⚠️ GROQ_API_KEY not set in environment")
        else:
            logger.info("✅ GROQ_API_KEY found")
        
        # Initialize API client
        logger.info("🌐 جاري تهيئة API Client...")
        api_client = DermaScanAPIClient(EXTERNAL_API_URL)
        set_api_client(api_client)
        
        # Check external API health
        try:
            if api_client.health_check():
                logger.info("✅ External API is healthy")
            else:
                logger.warning("⚠️ External API not reachable")
        except Exception as e:
            logger.error(f"❌ External API health check failed: {e}")
        
        # Build vectorstore and agent
        logger.info("🗄️ جاري بناء Vectorstore...")
        try:
            vectorstore = build_default_vectorstore()
            if vectorstore:
                logger.info("✅ Vectorstore loaded successfully")
                retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
            else:
                logger.warning("⚠️ Vectorstore not available")
                retriever = None
        except Exception as e:
            logger.error(f"❌ Failed to build vectorstore: {e}", exc_info=True)
            retriever = None
        
        # Build agents
        logger.info("🤖 جاري بناء Agent...")
        try:
            agent = build_agent(retriever, model_name=MODEL_NAME, temperature=TEMPERATURE)
            logger.info("✅ Primary agent built successfully")
        except Exception as e:
            logger.error(f"❌ Failed to build primary agent: {e}", exc_info=True)
            agent = None
        
        try:
            fallback_agent = build_agent(retriever, model_name=FALLBACK_MODEL_NAME, temperature=TEMPERATURE)
            logger.info("✅ Fallback agent built successfully")
        except Exception as e:
            logger.error(f"❌ Failed to build fallback agent: {e}", exc_info=True)
            fallback_agent = None
        
        logger.info("🎉 تم تهيئة النظام بنجاح!")
        
    except Exception as e:
        logger.error(f"❌ فشل تهيئة النظام: {e}", exc_info=True)
        raise

# FastAPI app with lifespan
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("🚀 FastAPI startup...")
    initialize_system()
    yield
    logger.info("🛑 FastAPI shutdown...")

app = FastAPI(title="DermaScan AI - RAG System", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models
class ChatRequest(BaseModel):
    thread_id: str
    message: str

class BookingRequest(BaseModel):
    thread_id: str
    city: str = ""

# FastAPI endpoints
@app.get("/health")
def health():
    return {
        "agent_ready": agent is not None,
        "api_client_ready": api_client is not None,
        "external_api_healthy": api_client.health_check() if api_client else False,
    }

@app.post("/chat")
def chat(req: ChatRequest):
    """Chat endpoint"""
    global agent
    
    if agent is None:
        return {"ok": False, "error": "Agent not initialized yet"}
    
    session = get_session(req.thread_id)
    set_session_storage(SESSIONS)
    set_current_thread_id(req.thread_id)
    
    try:
        last_analysis = session.get("last_analysis")
        analysis_ctx = ""
        if last_analysis:
            analysis_ctx = (
                f"\n\n[CURRENT IMAGE ANALYSIS — Classification: {last_analysis.get('label', 'N/A')}, "
                f"Confidence: {last_analysis.get('confidence_pct', 0):.1f}%, "
                f"Affected Area: {last_analysis.get('infection_pct', 0):.1f}%]"
            )
        
        full_prompt = req.message + analysis_ctx
        config = {"configurable": {"thread_id": req.thread_id}, "recursion_limit": 14}
        
        final_text = ""
        tool_calls = []
        
        for event in agent.stream(
            {"messages": [("user", full_prompt)]},
            config=config,
            stream_mode="values"
        ):
            last_msg = event["messages"][-1]
            if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
                for tc in last_msg.tool_calls:
                    tool_calls.append({"name": tc["name"], "args": tc["args"]})
            if last_msg.type == "ai" and last_msg.content:
                final_text = last_msg.content
        
        cleaned_text = re.sub(r'\[ANALYSIS_RESULT:\{.*?\}\]', '', final_text).strip()
        session["_turn_count"] = session.get("_turn_count", 0) + 1
        
        return {
            "ok": True,
            "reply": cleaned_text,
            "tool_calls": tool_calls,
        }
        
    except Exception as e:
        logger.error(f"Chat error: {e}", exc_info=True)
        return {"ok": False, "error": str(e)}

@app.post("/book-appointment")
def book_appointment(req: BookingRequest):
    """Book appointment via external API"""
    if api_client is None:
        return {"ok": False, "error": "API client not initialized"}
    
    result = api_client.book_appointment(req.thread_id, req.city)
    return result

@app.get("/image")
def get_image(name: str = "", thread_id: str = ""):
    """Get image from session"""
    try:
        session = get_session(thread_id) if thread_id else {}
        blobs = session.get("_image_blobs", {})
        
        if name not in blobs:
            return {"ok": False, "error": "Image not found"}
        
        return {"ok": True, "base64": blobs[name], "media_type": "image/png"}
    except Exception as e:
        return {"ok": False, "error": str(e)}

# Gradio Functions
def gradio_chat(message: str, history: list, thread_id: str = "default") -> str:
    """Gradio chat function"""
    if agent is None:
        return "System is not initialized yet. Please wait..."
    
    session = get_session(thread_id)
    set_session_storage(SESSIONS)
    set_current_thread_id(thread_id)
    
    try:
        last_analysis = session.get("last_analysis")
        analysis_ctx = ""
        if last_analysis:
            analysis_ctx = (
                f"\n\n[CURRENT IMAGE ANALYSIS — Classification: {last_analysis.get('label', 'N/A')}, "
                f"Confidence: {last_analysis.get('confidence_pct', 0):.1f}%, "
                f"Affected Area: {last_analysis.get('infection_pct', 0):.1f}%]"
            )
        
        full_prompt = message + analysis_ctx
        config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 14}
        
        final_text = ""
        for event in agent.stream(
            {"messages": [("user", full_prompt)]},
            config=config,
            stream_mode="values"
        ):
            last_msg = event["messages"][-1]
            if last_msg.type == "ai" and last_msg.content:
                final_text = last_msg.content
        
        cleaned_text = re.sub(r'\[ANALYSIS_RESULT:\{.*?\}\]', '', final_text).strip()
        session["_turn_count"] = session.get("_turn_count", 0) + 1
        
        return cleaned_text if cleaned_text else "No response generated."
        
    except Exception as e:
        logger.error(f"Gradio chat error: {e}", exc_info=True)
        return f"Error: {str(e)}"

def gradio_upload_image(image: Image.Image, thread_id: str = "default") -> tuple:
    """Handle image upload"""
    if api_client is None:
        return "API client not initialized", None, None, None
    
    try:
        result = api_client.upload_and_analyze_image(image, thread_id)
        
        if 'error' in result:
            return f"Error: {result['error']}", None, None, None
        
        session = get_session(thread_id)
        session["last_analysis"] = {
            "label": result.get("label"),
            "confidence_pct": result.get("confidence_pct"),
            "infection_pct": result.get("infection_pct"),
        }
        
        if 'images' in result:
            session["_image_blobs"].update(result['images'])
        
        analysis_text = f"""
✅ **تم تحليل الصورة بنجاح!**

**النتائج:**
- التصنيف: {result.get('label')}
- نسبة الثقة: {result.get('confidence_pct')}%
- المساحة المصابة: {result.get('infection_pct')}%

يمكنك الآن طرح أسئلة عن هذه الصورة.
"""
        
        orig_b64 = result.get('images', {}).get('original', '')
        mask_b64 = result.get('images', {}).get('mask', '')
        overlay_b64 = result.get('images', {}).get('overlay', '')
        
        orig_img = Image.open(io.BytesIO(base64.b64decode(orig_b64))) if orig_b64 else None
        mask_img = Image.open(io.BytesIO(base64.b64decode(mask_b64))) if mask_b64 else None
        overlay_img = Image.open(io.BytesIO(base64.b64decode(overlay_b64))) if overlay_b64 else None
        
        return analysis_text, orig_img, mask_img, overlay_img
        
    except Exception as e:
        logger.error(f"Upload error: {e}", exc_info=True)
        return f"Error: {str(e)}", None, None, None

# Build Gradio UI - ✅ الحل: ssr_mode=False لتعطيل Node.js
logger.info("🎨 جاري بناء واجهة Gradio...")
try:
    with gr.Blocks(title="DermaScan AI - Patient Portal", ssr_mode=False) as demo:
        gr.Markdown("""
        # 🏥 DermaScan AI - مساعد أمراض الجلدية
        
        ارفع صورة الجلد واحصل على تحليل بالذكاء الاصطناعي + اسأل أي أسئلة.
        """)
        
        thread_id_state = gr.State(value="default")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 رفع الصورة")
                image_input = gr.Image(type="pil", label="ارفع صورة الجلد")
                upload_btn = gr.Button("📎 ارفع وحلّل", variant="primary")
                analysis_output = gr.Markdown(label="نتيجة التحليل")
            
            with gr.Column(scale=2):
                gr.Markdown("### 💬 الدردشة مع DermaScan AI")
                chatbot = gr.Chatbot(label="المحادثة", height=400)
                msg_input = gr.Textbox(
                    label="رسالتك",
                    placeholder="اسأل عن حالتك الجلدية...",
                    lines=2
                )
                send_btn = gr.Button("🚀 إرسال", variant="primary")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 🖼️ الصور المحللة")
                orig_display = gr.Image(label="الصورة الأصلية", type="pil")
            with gr.Column():
                mask_display = gr.Image(label="قناع التجزئة", type="pil")
            with gr.Column():
                overlay_display = gr.Image(label="الصورة مع التغطية", type="pil")
        
        def user_message(user_msg, history):
            return "", history + [[user_msg, None]]
        
        def bot_response(history, thread_id):
            if not history:
                return history
            
            user_msg = history[-1][0]
            response = gradio_chat(user_msg, history[:-1], thread_id)
            
            history[-1][1] = response
            return history
        
        upload_btn.click(
            fn=gradio_upload_image,
            inputs=[image_input, thread_id_state],
            outputs=[analysis_output, orig_display, mask_display, overlay_display]
        )
        
        send_btn.click(
            fn=user_message,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot]
        ).then(
            fn=bot_response,
            inputs=[chatbot, thread_id_state],
            outputs=[chatbot]
        )
    
    logger.info("✅ تم بناء واجهة Gradio بنجاح")
    
except Exception as e:
    logger.error(f"❌ فشل بناء واجهة Gradio: {e}", exc_info=True)
    raise

logger.info("🎉 التطبيق جاهز للتشغيل!")

# ✅ Mount Gradio to FastAPI
app = gr.mount_gradio_app(app, demo, path="/")
logger.info("✅ تم ربط Gradio بـ FastAPI")

# ✅ تشغيل التطبيق - HF Spaces بيشغل تلقائياً
# لا حاجة لـ demo.launch() أو uvicorn.run()