"""
JARVIS — Image API Route
POST /api/image/analyze  — upload image + optional question → analysis + LLM response
POST /api/image/caption  — quick caption only
"""

import uuid
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from loguru import logger

from app.db.database import AsyncSession, get_db, create_session, get_session, add_message
from app.pipelines.image_pipeline import (
    analyze_image,
    generate_caption,
    load_image_from_bytes,
)
from app.pipelines.llm_pipeline import build_chat_prompt, generate_response
from app.memory.faiss_store import FAISSMemoryStore
from app.core.config import settings


router = APIRouter()


# ─────────────────────────────────────────────
# Response schemas
# ─────────────────────────────────────────────
class ImageAnalysisResponse(BaseModel):
    caption:          str
    detailed_caption: str
    top_categories:   list
    analysis:         str
    llm_response:     str
    session_id:       str
    size:             dict


# ─────────────────────────────────────────────
# Routes
# ─────────────────────────────────────────────
@router.post("/analyze", response_model=ImageAnalysisResponse)
async def analyze_image_endpoint(
    file:       UploadFile = File(...),
    question:   Optional[str] = Form(None),
    session_id: Optional[str] = Form(None),
    db:         AsyncSession = Depends(get_db),
):
    """
    Analyze an uploaded image using BLIP + CLIP.
    Optionally answer a question about the image using the LLM.
    """
    # Validate file type
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    # Read file
    contents = await file.read()
    if len(contents) > settings.MAX_UPLOAD_MB * 1024 * 1024:
        raise HTTPException(status_code=413, detail=f"Image too large (max {settings.MAX_UPLOAD_MB}MB)")

    # Load PIL image
    try:
        pil_image = load_image_from_bytes(contents)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid image: {str(e)}")

    # Run vision analysis
    try:
        vision_result = analyze_image(pil_image, user_question=question)
    except Exception as e:
        logger.error(f"Vision analysis error: {e}")
        raise HTTPException(status_code=500, detail=f"Vision analysis failed: {str(e)}")

    # Build LLM response
    image_caption = vision_result.get("caption", "")
    detailed = vision_result.get("detailed_caption", "")
    analysis = vision_result.get("analysis", "")

    user_question_text = question or "Describe this image in detail."

    prompt = build_chat_prompt(
        user_message=user_question_text,
        history=[],
        image_caption=f"{image_caption}. {detailed}",
        system_override=(
            "You are JARVIS, analyzing an image for the user. "
            "Use the provided image caption to give a thorough, helpful response. "
            "Be descriptive and insightful."
        ),
    )
    llm_response = generate_response(prompt, max_new_tokens=300)

    # Session handling
    session_id = session_id or str(uuid.uuid4())
    session = await get_session(db, session_id)
    if not session:
        await create_session(db, session_id, title=f"Image: {image_caption[:40]}")

    # Store in DB
    user_content = f"[IMAGE UPLOADED] {user_question_text}"
    await add_message(db, session_id, "user", user_content, modality="image")
    await add_message(db, session_id, "assistant", llm_response, modality="image")

    # Store in memory
    memory = FAISSMemoryStore.get_instance()
    memory.add(
        f"Image analysis: {image_caption}. {detailed}",
        session_id=session_id,
        metadata={"type": "image"},
    )

    return ImageAnalysisResponse(
        caption=image_caption,
        detailed_caption=detailed,
        top_categories=vision_result.get("top_categories", []),
        analysis=analysis,
        llm_response=llm_response,
        session_id=session_id,
        size=vision_result.get("size", {}),
    )


@router.post("/caption")
async def quick_caption(file: UploadFile = File(...)):
    """Quick image captioning — no LLM, just BLIP."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    contents = await file.read()
    try:
        pil_image = load_image_from_bytes(contents)
        caption = generate_caption(pil_image)
        return {"caption": caption, "width": pil_image.width, "height": pil_image.height}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))