app/workers/tasks.py

from app.workers.celery_app import celery_app
from app.core.database import SessionLocal
from app.models.user import User
from app.models.wardrobe import WardrobeItem
from app.models.body import BodyProfile
from app.models.tryon import TryOnTask
from app.services.storage import storage_service
from app.services.qdrant_client import qdrant_service
from app.services.tryon_service import tryon_service
from app.services.vision_service import vision_service
import uuid
import time
import logging
import random
import os
import io
import requests
from PIL import Image
from app.core.config import settings


# Configure logger
logger = logging.getLogger("celery_tasks")


@celery_app.task(name="app.workers.tasks.analyze_wardrobe_item")
def analyze_wardrobe_item(item_id: str):
    logger.info(f"Starting analysis for wardrobe item: {item_id}")
    db = SessionLocal()
    try:
        item_uuid = uuid.UUID(item_id) if isinstance(item_id, str) else item_id
        item = db.query(WardrobeItem).filter(WardrobeItem.id == item_uuid).first()
        if not item:
            logger.error(f"Item {item_id} not found in DB.")
            return False

        # Load image bytes (local /static file or remote URL)
        image_bytes = None
        local_path = None
        if item.image_url.startswith("http://localhost:") or "/static/" in item.image_url:
            parts = item.image_url.split("/static/")
            if len(parts) > 1:
                local_path = os.path.join(settings.UPLOAD_DIR, parts[1])

        try:
            if local_path and os.path.exists(local_path):
                with open(local_path, "rb") as f:
                    image_bytes = f.read()
            else:
                response = requests.get(item.image_url, timeout=15)
                response.raise_for_status()
                image_bytes = response.content
        except Exception as img_err:
            logger.error(f"Failed to load image from {item.image_url} / {local_path}: {img_err}")

        # Classify the garment via Gemini Vision API (no local ML).
        analysis = vision_service.classify_garment(image_bytes) if image_bytes else None

        # Fallback to simulated metadata if the API is unavailable / failed.
        if analysis is None:
            time.sleep(0.5)
            cat = random.choice(["shirt", "pants", "dress", "jacket", "skirt"])
            analysis = {
                "category": cat,
                "color_hex": f"#{random.randint(0, 0xFFFFFF):06x}",
                "style_tags": ["casual", "minimalist"] if cat in ["shirt", "pants"] else ["office", "formal"],
            }
            logger.info(f"Wardrobe item {item_id} analyzed with simulated metadata.")
        else:
            logger.info(f"Wardrobe item {item_id} classified via Gemini Vision: {analysis['category']}.")

        # Store a mock 512-d vector in Qdrant (embeddings are not computed locally).
        vector_id = uuid.uuid4()
        mock_vector = [random.uniform(-0.1, 0.1) for _ in range(512)]
        qdrant_service.upsert_vector(vector_id, mock_vector, {
            "user_id": str(item.user_id),
            "category": analysis["category"],
            "style_tags": analysis["style_tags"],
        })

        item.category = analysis["category"]
        item.color_hex = analysis["color_hex"]
        item.style_tags = analysis["style_tags"]
        item.qdrant_vector_id = vector_id
        db.commit()
        return True

    except Exception as e:
        db.rollback()
        logger.error(f"Error analyzing wardrobe item {item_id}: {str(e)}", exc_info=True)
        return False
    finally:
        db.close()

@celery_app.task(name="app.workers.tasks.reconstruct_body_mesh")
def reconstruct_body_mesh(profile_id: str):
    logger.info(f"Starting 3D body reconstruction for profile: {profile_id}")
    db = SessionLocal()
    try:
        profile_uuid = uuid.UUID(profile_id) if isinstance(profile_id, str) else profile_id
        profile = db.query(BodyProfile).filter(BodyProfile.id == profile_uuid).first()
        if not profile:
            logger.error(f"BodyProfile {profile_id} not found.")
            return False

        # Simulate MediaPipe Keypoint extraction & SMPL reconstruction
        time.sleep(3.0)
        
        # Mock keypoints and SMPL betas parameters
        mock_keypoints = {
            "shoulder_left": [0.35, 0.25, -0.05],
            "shoulder_right": [0.65, 0.25, -0.05],
            "hip_left": [0.40, 0.55, 0.0],
            "hip_right": [0.60, 0.55, 0.0]
        }
        mock_shape_params = [random.uniform(-1.5, 1.5) for _ in range(10)]
        
        # Mock GLB avatar file upload to R2
        mock_glb_content = b"Mock GLB file bytes representing 3D avatar mesh"
        mesh_path = f"users/{profile.user_id}/body/avatar.glb"
        avatar_url = storage_service.upload_file(mock_glb_content, mesh_path, "model/gltf-binary")

        if avatar_url:
            profile.keypoints_2d = mock_keypoints
            profile.shape_params = mock_shape_params
            profile.avatar_mesh_url = avatar_url
            db.commit()
            logger.info(f"Successfully reconstructed body for profile {profile_id}. GLB url: {avatar_url}")
            return True
        else:
            logger.error("Failed to upload GLB mesh to R2.")
            return False
            
    except Exception as e:
        db.rollback()
        logger.error(f"Error reconstructing body mesh {profile_id}: {str(e)}", exc_info=True)
        return False
    finally:
        db.close()


@celery_app.task(name="app.workers.tasks.run_tryon")
def run_tryon(task_id: str):
    logger.info(f"Starting Virtual Try-On using provider '{settings.TRYON_PROVIDER}' for task: {task_id}")
    db = SessionLocal()
    try:
        task_uuid = uuid.UUID(task_id) if isinstance(task_id, str) else task_id
        task = db.query(TryOnTask).filter(TryOnTask.id == task_uuid).first()
        if not task:
            logger.error(f"TryOnTask {task_id} not found.")
            return False

        task.status = "PROCESSING"
        db.commit()

        # 1. Load person image
        person_img = None
        if task.person_image_url:
            local_person_path = None
            if task.person_image_url.startswith("http://localhost:") or "/static/" in task.person_image_url:
                parts = task.person_image_url.split("/static/")
                if len(parts) > 1:
                    local_person_path = os.path.join(settings.UPLOAD_DIR, parts[1])
            try:
                if local_person_path and os.path.exists(local_person_path):
                    person_img = Image.open(local_person_path).convert("RGB")
                else:
                    response = requests.get(task.person_image_url, timeout=15)
                    person_img = Image.open(io.BytesIO(response.content)).convert("RGB")
            except Exception as e:
                logger.warning(f"Could not load custom person image from {task.person_image_url}: {e}")
        
        # Fallback to default model/hero photo
        if person_img is None:
            default_hero_paths = [
                "/home/llm/MinhPV/AI_Virtual_Wardrobe/frontend/src/assets/hero.jpg",
                "/home/llm/MinhPV/AI_Virtual_Wardrobe/src/assets/hero.jpg",
            ]
            for path in default_hero_paths:
                if os.path.exists(path):
                    try:
                        person_img = Image.open(path).convert("RGB")
                        break
                    except Exception as e:
                        logger.error(f"Error loading hero from {path}: {e}")
            if person_img is None:
                logger.warning("No hero.jpg found on disk. Creating a blank image for try-on simulation.")
                person_img = Image.new("RGB", (512, 512), (240, 240, 240))

        # 2. Load garments
        from app.models.wardrobe import WardrobeItem
        garments = []
        if task.garment_item_ids:
            garments = db.query(WardrobeItem).filter(WardrobeItem.id.in_(task.garment_item_ids)).all()
        elif task.garment_item_id:
            garment = db.query(WardrobeItem).filter(WardrobeItem.id == task.garment_item_id).first()
            if garment:
                garments = [garment]
                
        if not garments:
            raise ValueError("No valid garments found for the try-on task.")

        # 3. Collect garment images (bytes + category + url) for the generation service
        garment_inputs = []  # list[(bytes, category)]
        garment_urls = []    # aligned URLs, used by URL-based providers (e.g. Replicate)
        for garment in garments:
            logger.info(f"Loading garment item ID {garment.id} (category: {garment.category})")

            garment_bytes = None
            if garment.image_url:
                local_garment_path = None
                if garment.image_url.startswith("http://localhost:") or "/static/" in garment.image_url:
                    parts = garment.image_url.split("/static/")
                    if len(parts) > 1:
                        local_garment_path = os.path.join(settings.UPLOAD_DIR, parts[1])
                try:
                    if local_garment_path and os.path.exists(local_garment_path):
                        with open(local_garment_path, "rb") as f:
                            garment_bytes = f.read()
                    else:
                        response = requests.get(garment.image_url, timeout=15)
                        response.raise_for_status()
                        garment_bytes = response.content
                except Exception as e:
                    logger.error(f"Could not load garment image from {garment.image_url}: {e}")
                    raise ValueError(f"Không thể tải ảnh sản phẩm {garment.image_url}: {e}")

            if garment_bytes is None:
                raise ValueError(f"Không có ảnh sản phẩm hợp lệ cho item {garment.id}")

            garment_inputs.append((garment_bytes, (garment.category or "clothing").lower()))
            garment_urls.append(garment.image_url)

        # 4. Encode the person image and generate the try-on via the configured provider
        person_bio = io.BytesIO()
        person_img.convert("RGB").save(person_bio, format="JPEG", quality=95)
        person_bytes = person_bio.getvalue()

        result_image_bytes = tryon_service.generate(
            person_bytes=person_bytes,
            garments=garment_inputs,
            person_url=task.person_image_url,
            garment_urls=garment_urls,
        )
        if not result_image_bytes:
            raise ValueError("Dịch vụ tạo ảnh thử đồ không trả về ảnh.")

        # Normalize provider output to JPEG so the .jpg path and content-type stay consistent
        out_img = Image.open(io.BytesIO(result_image_bytes)).convert("RGB")
        out_bio = io.BytesIO()
        out_img.save(out_bio, format="JPEG", quality=95)
        result_image_bytes = out_bio.getvalue()
        logger.info(f"Successfully generated try-on image via provider '{settings.TRYON_PROVIDER}'.")

        # 5. Upload final image
        result_path = f"users/{task.user_id}/try-on/{task.id}.jpg"
        result_url = storage_service.upload_file(result_image_bytes, result_path, "image/jpeg")

        if result_url:
            task.status = "SUCCESS"
            task.result_url = result_url
            db.commit()
            logger.info(f"Successfully completed Try-On task {task_id}. Image url: {result_url}")
            return True
        else:
            task.status = "FAILED"
            task.error_message = "Không thể lưu trữ kết quả thử đồ lên hệ thống lưu trữ local/R2."
            db.commit()
            return False
            
    except Exception as e:
        db.rollback()
        logger.error(f"Error in Virtual Try-On task {task_id}: {str(e)}", exc_info=True)
        if task:
            err = str(e)
            low = err.lower()
            if "429" in err or "quota" in low or "resource_exhausted" in low or "limit: 0" in low:
                friendly = (
                    "Đã hết quota tạo ảnh của nhà cung cấp (model sinh ảnh không khả dụng "
                    "trên gói miễn phí). Vui lòng bật billing hoặc đổi provider rồi thử lại."
                )
            else:
                friendly = f"Tạo ảnh thử đồ thất bại: {err}"
            task.status = "FAILED"
            task.error_message = friendly
            db.commit()
        return False
    finally:
        db.close()