import os
import io
import uuid
import shutil
import torch
import numpy as np
import asyncio
import logging
import time
from PIL import Image
from pathlib import Path
from typing import List, Dict
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
import ollama
from transformers import AutoModel

# --- Config & Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
BASE_DIR = Path(__file__).parent.resolve()
UPLOADS_DIR = BASE_DIR / "uploads"
STATIC_DIR = BASE_DIR / "static"
MAX_FILE_SIZE = 20 * 1024 * 1024  # 20MB
ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
SESSION_EXPIRY = 3600  # 1 hour

OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "qwen2.5vl:3b")
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://localhost:11434")
# กำหนด timeout เป็น 300.0 วินาที (5 นาที) สำหรับกรณีรันบน CPU หรือระบบที่มีทรัพยากรจำกัด
ollama_client = ollama.AsyncClient(host=OLLAMA_HOST, timeout=300.0)

# Magic Bytes for file validation
MAGIC_BYTES = {
    b"\x89PNG": "png",
    b"\xff\xd8\xff": "jpeg",
    b"RIFF": "webp",  # WebP files start with RIFF...WEBP
}

# Ensure directories exist
UPLOADS_DIR.mkdir(exist_ok=True)
STATIC_DIR.mkdir(exist_ok=True)

# Global models
models: Dict = {}


async def periodic_cleanup_task():
    """Background task to clean up expired sessions regularly."""
    while True:
        try:
            logger.info("🧹 Starting periodic session cleanup...")
            now = time.time()
            if UPLOADS_DIR.exists():
                for item in UPLOADS_DIR.iterdir():
                    if item.is_dir():
                        # Check directory modification time
                        mtime = item.stat().st_mtime
                        if now - mtime > SESSION_EXPIRY:
                            logger.info(f"🧹 Cleaning up expired session: {item.name}")
                            # Delete in a thread pool to avoid blocking the event loop
                            await asyncio.to_thread(shutil.rmtree, item)
        except Exception as e:
            logger.error(f"❌ Error in periodic cleanup task: {e}")

        # Run cleanup every 10 minutes
        await asyncio.sleep(600)


@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- Load Models ---
    try:
        logger.info("🚀 กำลังโหลดโมเดล MAGI...")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"💻 ใช้ฮาร์ดแวร์อุปกรณ์: {device}")

        models["magi"] = AutoModel.from_pretrained(
            "ragavsachdeva/magi", trust_remote_code=True
        ).to(device)
        models["magi"].eval()
        logger.info("✅ โหลดโมเดล MAGI สำเร็จ")
    except Exception as e:
        logger.error(f"❌ Failed to load models: {e}")
        # โยน Exception เพื่อให้โปรแกรม Crash-on-startup (Fail-fast) จะได้รู้ว่าโมเดลโหลดไม่ผ่าน
        raise RuntimeError(f"Could not load required models: {e}")

    # --- ตรวจสอบสถานะการเชื่อมต่อของ Ollama ---
    try:
        logger.info(f"🔍 กำลังตรวจสอบการเชื่อมต่อกับ Ollama Server ที่ {OLLAMA_HOST}...")
        # ใช้ timeout สั้นๆ สำหรับขั้นตอน Startup เพื่อไม่ให้โปรแกรมค้างหาก Ollama ยังไม่ได้เปิด
        temp_client = ollama.AsyncClient(host=OLLAMA_HOST, timeout=3.0)
        models_list = await temp_client.list()
        logger.info("✅ เชื่อมต่อกับ Ollama Server สำเร็จ!")
        
        # ตรวจสอบว่าโมเดลพร้อมใช้งานหรือไม่
        available_models = [m.get("name", m.get("model", "")) for m in models_list.get("models", [])]
        logger.info(f"📦 โมเดลที่มีใน Ollama: {available_models}")
        
        if OLLAMA_MODEL not in available_models and f"{OLLAMA_MODEL}:latest" not in available_models:
            logger.warning(
                f"⚠️ ไม่พบโมเดล '{OLLAMA_MODEL}' ใน Ollama Server! "
                f"กรุณารันคำสั่ง 'ollama pull {OLLAMA_MODEL}' ใน Terminal ก่อนเริ่มวิเคราะห์"
            )
    except Exception as e:
        logger.error(
            f"❌ ไม่สามารถเชื่อมต่อกับ Ollama Server ได้ ({e}) "
            f"โปรดตรวจสอบว่า Ollama Server กำลังรันอยู่บนโฮสต์ {OLLAMA_HOST} และโมเดลติดตั้งไว้แล้ว"
        )

    # Start periodic cleanup as a background task
    cleanup_task = asyncio.create_task(periodic_cleanup_task())

    yield

    # Cancel background cleanup
    cleanup_task.cancel()
    try:
        await cleanup_task
    except asyncio.CancelledError:
        pass

    # --- Cleanup ---
    logger.info("🧹 Cleaning up models...")
    models.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def get_ready_models():
    """Dependency to ensure models are loaded."""
    if not models.get("magi"):
        raise HTTPException(status_code=503, detail="Models not ready")
    return models


app = FastAPI(lifespan=lifespan)

# Mount static
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/uploads/{session_id}/{filename}")
async def serve_upload(session_id: str, filename: str):
    try:
        uuid.UUID(session_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid session ID")

    # Sanitize filename to prevent Path Traversal
    filename = Path(filename).name
    file_path = (UPLOADS_DIR / session_id / filename).resolve()

    try:
        # Verify that the path is actually inside UPLOADS_DIR
        file_path.relative_to(UPLOADS_DIR.resolve())
    except ValueError:
        raise HTTPException(status_code=403, detail="Access denied")

    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found")

    return FileResponse(file_path)


def validate_magic_bytes(content: bytes):
    """Verify file type using magic bytes."""
    for magic in MAGIC_BYTES.keys():
        if content.startswith(magic):
            return True
    raise HTTPException(
        status_code=400,
        detail="Invalid file content: Magic bytes do not match allowed image types.",
    )


def validate_upload_meta(file: UploadFile):
    ext = Path(file.filename or "").suffix.lower()
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: {ext}. Allowed: {ALLOWED_EXTENSIONS}",
        )


async def get_qwen_description(image_crop: Image.Image, panel_index: int) -> str:
    # ขยายภาพ 2 เท่า ด้วย LANCZOS เพื่อรักษาความคมชัดของภาพมังงะและตัวอักษรขนาดเล็ก
    width, height = image_crop.size

    # ดึงการปรับขนาดรูปแบบ PIL ออกไปทำนอก Thread หลักเพื่อป้องกันความเร็วตก
    upscaled_image = await asyncio.to_thread(
        image_crop.resize, (width * 2, height * 2), Image.LANCZOS
    )

    img_byte_arr = io.BytesIO()
    await asyncio.to_thread(upscaled_image.save, img_byte_arr, format="PNG")
    img_bytes = img_byte_arr.getvalue()

    prompt = f"""นี่คือกรอบมังงะช่องที่ {panel_index} 
    จงเขียน 'บทอ่านออกเสียง' (Script for Reading) สำหรับเนื้อหาในภาพนี้ โดยมีเงื่อนไขดังนี้:
    1. ให้บรรยายสถานการณ์และบทสนทนาให้เป็นเนื้อความเดียวกันแบบนิยาย
    2. ไม่ต้องมีหัวข้อ (เช่น ไม่ต้องมี 1. 2. 3. หรือ คำว่า 'สถานการณ์:')
    3. ไม่ต้องมีสัญลักษณ์พิเศษ หรือเครื่องหมายคำพูดที่ซับซ้อน
    4. ใช้ภาษาไทยที่สละสลวย อ่านแล้วเข้าใจทันทีว่าใครพูดอะไร หรือเกิดอะไรขึ้น
    5. เน้นเนื้อหาที่จะนำไปใช้กับระบบ Text-to-Speech (อ่านออกเสียง) ได้ทันที
    """

    logger.info(f"🤖 [Panel {panel_index}] กำลังส่งคำร้องไปยัง Ollama ({OLLAMA_MODEL})...")
    start_time = time.time()
    try:
        response = await ollama_client.chat(
            model=OLLAMA_MODEL,
            messages=[{"role": "user", "content": prompt, "images": [img_bytes]}],
            options={
                "num_predict": 250,  # จำกัดจำนวน token คำตอบ เพื่อเร่งความเร็วบน CPU
                "temperature": 0.2,
            }
        )
        elapsed = time.time() - start_time
        logger.info(f"✅ [Panel {panel_index}] ได้รับคำอธิบายจาก Ollama แล้ว (ใช้เวลา: {elapsed:.2f} วินาที)")
        return response["message"]["content"].strip()
    except Exception as e:
        elapsed = time.time() - start_time
        logger.error(f"❌ [Panel {panel_index}] เกิดข้อผิดพลาดในการเชื่อมต่อ Ollama (หลัง {elapsed:.2f} วินาที): {e}")
        return f"Error ในการบรรยาย: {e}"


@app.get("/")
async def read_index():
    index_path = STATIC_DIR / "index.html"
    if not index_path.exists():
        return HTMLResponse(content="<h1>index.html not found</h1>", status_code=404)
    # Read file inside a thread pool
    content = await asyncio.to_thread(index_path.read_text, encoding="utf-8")
    return HTMLResponse(content=content)


@app.post("/analyze")
async def analyze_manga(
    file: UploadFile = File(...), loaded_models: Dict = Depends(get_ready_models)
):
    validate_upload_meta(file)
    logger.info(f"📥 เริ่มต้นกระบวนการวิเคราะห์ไฟล์: {file.filename}")
    start_time = time.time()
    try:
        # 1. Read file in chunks to prevent OOM
        content_length = file.headers.get("content-length")
        if content_length and int(content_length) > MAX_FILE_SIZE:
            logger.warning(f"❌ ไฟล์มีขนาดใหญ่เกินกำหนด: {content_length} bytes")
            raise HTTPException(status_code=413, detail="File too large. Max 20MB.")

        chunks = []
        bytes_read = 0
        while True:
            chunk = await file.read(1024 * 1024)  # 1MB chunk
            if not chunk:
                break
            bytes_read += len(chunk)
            if bytes_read > MAX_FILE_SIZE:
                logger.warning(f"❌ ไฟล์มีขนาดใหญ่เกินกำหนดระหว่างอ่านข้อมูล: {bytes_read} bytes")
                raise HTTPException(status_code=413, detail="File too large. Max 20MB.")
            chunks.append(chunk)
        content = b"".join(chunks)

        # 2. Validate Magic Bytes
        validate_magic_bytes(content)

        session_id = str(uuid.uuid4())
        session_dir = UPLOADS_DIR / session_id
        logger.info(f"📁 สร้าง Session ID: {session_id}")

        # Create session directory
        await asyncio.to_thread(session_dir.mkdir, parents=True, exist_ok=True)

        file_path = session_dir / "original.png"
        # Write bytes in thread pool
        await asyncio.to_thread(file_path.write_bytes, content)
        logger.info(f"💾 บันทึกรูปต้นฉบับสำเร็จที่: {file_path}")

        # 3. Run MAGI Panel Detection
        logger.info("🧠 กำลังวิเคราะห์เลย์เอาต์มังงะและตรวจจับช่อง (MAGI Panel Detection)...")
        pil_img = await asyncio.to_thread(Image.open, io.BytesIO(content))
        pil_img = pil_img.convert("RGB")
        img_np = np.array(pil_img)

        magi_model = loaded_models["magi"]
        loop = asyncio.get_running_loop()

        def run_magi():
            with torch.no_grad():
                return magi_model.predict_detections_and_associations([img_np])

        magi_start = time.time()
        results = await loop.run_in_executor(None, run_magi)
        res = results[0]
        panels = res.get("panels", [])
        logger.info(f"✅ ตรวจจับเสร็จสิ้น! พบกรอบมังงะทั้งหมด {len(panels)} กรอบ (ใช้เวลา: {time.time() - magi_start:.2f} วินาที)")

        if not panels:
            logger.warning("⚠️ ไม่พบช่อง/กรอบมังงะในหน้ากระดาษนี้")

        # 4. Prepare and Crop Panels
        output_panels = []
        crop_tasks = []

        async def process_panel(i, box):
            x1, y1, x2, y2 = [int(val) for val in box]
            pad = 5

            # Crop image
            crop = pil_img.crop(
                (
                    max(0, x1 - pad),
                    max(0, y1 - pad),
                    min(pil_img.width, x2 + pad),
                    min(pil_img.height, y2 + pad),
                )
            )

            # Save crop in thread pool
            panel_filename = f"panel_{i+1}.png"
            panel_path = session_dir / panel_filename
            await asyncio.to_thread(crop.save, panel_path)
            return crop, panel_filename

        # Crop all panels
        for i, box in enumerate(panels):
            crop_tasks.append(process_panel(i, box))

        cropped_results = await asyncio.gather(*crop_tasks)
        logger.info(f"✂️ ดำเนินการตัดและจัดเก็บรูปภาพย่อยจำนวน {len(cropped_results)} ภาพเรียบร้อย")

        # 5. Call Ollama sequentially/in parallel (using Semaphore to limit concurrency)
        # ปรับเหลือ 1 เพื่อรันทีละกรอบ ป้องกันการแย่งทรัพยากร CPU จนเกิดอาการ Timeout หรือดับ
        sem = asyncio.Semaphore(1)

        async def get_description_with_limit(crop, index):
            async with sem:
                return await get_qwen_description(crop, index)

        logger.info(f"📢 กำลังส่งทั้ง {len(cropped_results)} ช่องการ์ตูนไปให้ Ollama บรรยายภาพ (รันทีละ 1 ช่องเพื่อป้องกัน CPU เกินพิกัด)...")
        desc_tasks = [
            get_description_with_limit(crop, i + 1)
            for i, (crop, _) in enumerate(cropped_results)
        ]
        descriptions = await asyncio.gather(*desc_tasks)

        full_transcript = []
        for i, (crop, panel_filename) in enumerate(cropped_results):
            description = descriptions[i]
            full_transcript.append(f"--- กรอบที่ {i+1} ---\n{description}\n")
            output_panels.append(
                {
                    "index": i + 1,
                    "image_url": f"/uploads/{session_id}/{panel_filename}",
                    "description": description,
                }
            )

        # 6. Save Detection Visualization
        logger.info("🎨 กำลังวาดกรอบเลย์เอาต์มังงะและสร้างรูปภาพภาพรวม...")
        detection_result_path = session_dir / "detection_result.png"
        await asyncio.to_thread(
            magi_model.visualise_single_image_prediction,
            img_np,
            res,
            filename=str(detection_result_path),
        )

        # 7. Save Report
        report_path = session_dir / "report.txt"
        report_content = (
            f"### รายงานการวิเคราะห์ (Session: {session_id}) ###\n"
            f"จำนวนกรอบทั้งหมด: {len(panels)}\n"
            + "=" * 50
            + "\n\n"
            + "".join(full_transcript)
        )
        await asyncio.to_thread(
            report_path.write_text, report_content, encoding="utf-8"
        )
        
        total_time = time.time() - start_time
        logger.info(f"✨ เสร็จสมบูรณ์ทุกขั้นตอนเรียบร้อย! (เวลาทั้งหมด: {total_time:.2f} วินาที)")

        return {
            "session_id": session_id,
            "detection_result_url": f"/uploads/{session_id}/detection_result.png",
            "panels": output_panels,
        }

    except HTTPException as e:
        logger.error(f"❌ HTTP Exception: {e.detail}")
        raise
    except Exception as e:
        logger.exception("❌ เกิดข้อผิดพลาดร้ายแรงระหว่างกระบวนการวิเคราะห์ภาพ:")
        return JSONResponse(status_code=500, content={"error": "Internal server error"})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)