import os
import io
import uvicorn
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

from chatbot_updated import (
    chatbot_updated,
    detect_artifact,
    text_to_speech,
    cleanup_audio_file,
)

app = FastAPI(title="Egyptian Artifact Chatbot")

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

# =========================
# GET /
# =========================
@app.get("/")
async def root():
    return {
        "status": "ok",
        "message": "Egyptian Artifact Chatbot API is running!",
        "endpoints": {
            "POST /chat": "Send question + optional image/audio",
            "GET /health": "Health check",
            "GET /audio/{filename}": "Get TTS audio file",
            "GET /docs": "Swagger UI"
        }
    }

# =========================
# GET /health
# =========================
@app.get("/health")
async def health():
    return {"status": "ok"}

# =========================
# POST /chat
# Response:
#   answer        → نص الإجابة
#   detected      → اسم الأثر اللي YOLO اكتشفه
#   annotated_img → الصورة بعد الـ detection (base64 JPEG)
#   audio_url     → رابط ملف الصوت
# =========================
@app.post("/chat")
async def chat(
    question: str        = Form(default=""),
    image:    UploadFile = File(default=None),
    audio:    UploadFile = File(default=None),
):
    # ============================
    # 1. Input: صوت > نص
    # ============================
    input_question = question.strip() if question else ""

    if audio is not None:
        audio_bytes = await audio.read()
        if audio_bytes:
            input_question = audio_bytes

    # ============================
    # 2. الصورة — Detection
    # ============================
    img           = None
    detected_name = None
    annotated_b64 = None

    if image is not None:
        image_bytes = await image.read()
        if image_bytes:
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

            # Detection
            detected_name, annotated_img = detect_artifact(img)

            # نحفظ الصورة المحددة في temp folder
            import tempfile, uuid
            annotated_filename = f"annotated_{uuid.uuid4().hex}.jpg"
            annotated_path = os.path.join(tempfile.gettempdir(), annotated_filename)
            annotated_img.save(annotated_path, format="JPEG")
            annotated_b64 = annotated_filename  # بنحفظ الاسم بس

    # ============================
    # 3. صورة بس من غير سؤال → detection فقط
    # ============================
    if not input_question:
        if img is not None:
            return JSONResponse({
                "answer":        None,
                "detected":      detected_name,
                "annotated_url": f"/image/{annotated_b64}" if annotated_b64 else None,
                "audio_url":     None,
            })
        else:
            return JSONResponse(
                status_code=400,
                content={"error": "Please provide a question, audio, or image."}
            )

    # ============================
    # 4. الشات بوت
    # ============================
    answer = str(chatbot_updated(input_question, image=img)).strip()

    # ============================
    # 5. TTS
    # ============================
    audio_url  = None
    audio_file = text_to_speech(answer)
    if audio_file and os.path.exists(audio_file):
        audio_url = f"/audio/{os.path.basename(audio_file)}"

    # ============================
    # 6. Response
    # ============================
    return JSONResponse({
        "answer":        answer,
        "detected":      detected_name,
        "annotated_url": f"/image/{annotated_b64}" if annotated_b64 else None,
        "audio_url":     audio_url,
    })

# =========================
# GET /audio/{filename}
# =========================
@app.get("/audio/{filename}")
async def get_audio(filename: str):
    import tempfile
    path = os.path.join(tempfile.gettempdir(), filename)
    if not os.path.exists(path):
        return JSONResponse(status_code=404, content={"error": "File not found."})
    return FileResponse(path, media_type="audio/mpeg", filename=filename)

# =========================
# GET /image/{filename}
# بيرجع الصورة بعد الـ detection
# =========================
@app.get("/image/{filename}")
async def get_image(filename: str):
    import tempfile
    path = os.path.join(tempfile.gettempdir(), filename)
    if not os.path.exists(path):
        return JSONResponse(status_code=404, content={"error": "Image not found."})
    return FileResponse(path, media_type="image/jpeg", filename=filename)

# =========================
# POST /detect
# يرجع الصورة بعد الـ detection مباشرة كـ file
# =========================
@app.post("/detect")
async def detect(image: UploadFile = File(...)):
    image_bytes = await image.read()
    if not image_bytes:
        return JSONResponse(status_code=400, content={"error": "No image provided."})

    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    detected_name, annotated_img = detect_artifact(img)

    # نرجع الصورة مباشرة كـ JPEG
    buf = io.BytesIO()
    annotated_img.save(buf, format="JPEG")
    buf.seek(0)

    from fastapi.responses import StreamingResponse
    headers = {"X-Detected": detected_name or "none"}
    return StreamingResponse(buf, media_type="image/jpeg", headers=headers)

# =========================
# Run
# =========================
if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)