from pathlib import Path
from typing import Optional

from fastapi import HTTPException, status
from fastapi.concurrency import run_in_threadpool
from sqlmodel import Session

from app.core.logging import setup_logger
from app.core.constants import ERR_UNEXPECTED
from app.models.product_template import ProductTemplate
from app.models.project_master import ProjectMaster
from app.models.product_image import ProductImage
from app.repositories.product_template_repo import ProductTemplateRepository
from app.repositories.project_repo import ProjectRepository
from app.repositories.product_image_repo import ProductImageRepository
from app.utils.image import (
    validate_overlay_file,
    create_thumbnail_from_path,
    resize_image_to_size,
    save_upload_file_streamed,
    safe_delete,
    apply_color_layer_to_image,
)
from app.utils.generate_final_from_mask import generate_final_image_from_mask
from app.utils.paths import (
    ASSETS_ROOT,
    final_dir,
    final_thumb_dir,
    overlay_dir,
    rel_path,
    abs_from_db,
)
from app.utils.image_crop import crop_transparent_canvas
from app.utils.image_background import normalize_white_background
from app.utils.colorize import colorize_preview_with_mask

logger = setup_logger(__name__)

class ProductImageService:
    """
    Handles generation and regeneration of final product images.
    Files are stored under per-project final_images directories.
    """

    @staticmethod
    async def generate_for_project(
        session: Session,
        project_id: int,
        template_unique_name: str,
        current_user_id: Optional[int] = None,
    ) -> ProductImage:
        """
        Generate final image(s) for a project using the given template.
        Resulting files are stored in:
            assets/projects/project_[id]/final_images/
        DB paths stored as 'assets/...' relative strings.
        """
        # 1. Template lookup
        tpl: Optional[ProductTemplate] = ProductTemplateRepository.get_by_unique_name(session, template_unique_name)
        if not tpl:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Template not found")

        # 2. Project lookup
        project: Optional[ProjectMaster] = ProjectRepository.get_project(session, project_id)
        if not project or project.is_deleted:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Project not found")

        # 3. Ensure project has overlay
        if not project.overlay_image_path:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Project has no overlay image")

        # 4. Resolve absolute paths for PSD (template) and overlay/logo using abs_from_db
        tpl_psd_abs = abs_from_db(tpl.file_path) or Path(tpl.file_path)
        overlay_abs = abs_from_db(project.overlay_image_path) or Path(project.overlay_image_path)
        logo_abs = abs_from_db(project.logo_image_path) if project.logo_image_path else None

        # 5. Ensure final dirs exist
        out_final_dir = final_dir(project_id)
        out_final_dir.mkdir(parents=True, exist_ok=True)
        out_final_thumb_dir = final_thumb_dir(project_id)
        out_final_thumb_dir.mkdir(parents=True, exist_ok=True)

        # 6. Call generate_final_image_from_mask (CPU/IO-bound) in threadpool
        try:
            mask_image = tpl.main_black_mask_path if any(w in tpl.unique_name for w in ("mug", "cup")) else tpl.main_white_mask_path
            final_path_str = await run_in_threadpool(
                generate_final_image_from_mask,
                tpl.base_image_path,
                tpl.main_white_mask_path,
                str(overlay_abs),
                str(out_final_dir),
                tpl.unique_name
            )
        except HTTPException:
            raise
        except Exception:
            logger.exception("Failed to generate final image for template %s and project %s", tpl.id, project_id)
            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to generate final image")

        final_abs = Path(final_path_str)

        # 7. Create thumbnail in project's final_images/thumbnail
        try:
            final_thumb_abs = create_thumbnail_from_path(final_abs, out_final_thumb_dir)
        except Exception:
            # cleanup final if thumbnail generation fails (design choice)
            safe_delete(final_abs)
            logger.exception("Failed to create final thumbnail for %s", final_abs)
            raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to generate final thumbnail")

        # 8. Optional product-resized copy
        product_resized_path: Optional[Path] = None
        width = tpl.width
        height = tpl.height
        if width and height:
            try:
                w = int(width)
                h = int(height)
                product_resized_path = resize_image_to_size(final_abs, out_final_dir, w, h)
            except Exception:
                logger.exception("Failed to create product-resized image for %s", tpl.unique_name)

        # 9. Persist ProductImage record (store DB-friendly 'assets/...' paths)
        def _db_rel(p: Path | None) -> Optional[str]:
            return rel_path(p)

        product_image = ProductImage(
            project_master_id=project.id,
            product_template_id=tpl.id,
            final_image_path=_db_rel(final_abs),
            final_image_thumbnail_path=_db_rel(final_thumb_abs),
            final_image_product_path=_db_rel(product_resized_path) if product_resized_path else None,
        )

        product_image = ProductImageRepository.create(session, product_image)
        return product_image

    @staticmethod
    async def regenerate_single_image(
        db: Session,
        product_image_id: int,
        new_overlay_file,
        color: str | None,
        current_user,
    ):
        """
        Replace overlay (temporary), regenerate final image and derivatives,
        update DB atomically, and cleanup old files on success.
        """
        # Validate overlay payload
        try:
            validate_overlay_file(new_overlay_file)
        except Exception:
            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=ERR_UNEXPECTED)

        # Fetch ProductImage row
        img_row = db.get(ProductImage, product_image_id)
        if not img_row or img_row.is_deleted:
            raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Image record not found")

        # Fetch related objects
        template = ProductTemplateRepository.get_by_id(db, img_row.product_template_id)
        if not template:
            raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Associated template not found")

        project = ProjectRepository.get_project(db, img_row.project_master_id)
        if not project or project.is_deleted:
            raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Project not found")

        project_id = project.id

        # Resolve PSD path
        tpl_psd_abs = abs_from_db(template.file_path) or Path(template.file_path)

        # Save uploaded overlay to project's overlay dir (temporary)
        try:
            overlay_dir(project_id).mkdir(parents=True, exist_ok=True)
            overlay_temp_abs = (await save_upload_file_streamed(new_overlay_file, overlay_dir(project_id))).resolve()
            crop_transparent_canvas(overlay_temp_abs)
            normalize_white_background(overlay_temp_abs, target_rgb=(207, 207, 207))
        except Exception:
            logger.exception("Failed to store temporary overlay for project_id=%s", project_id)
            raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to store temporary overlay image")

        # Logo absolute path if present
        logo_abs = abs_from_db(project.logo_image_path) if project.logo_image_path else None

        # Ensure final dirs exist
        out_final_dir = final_dir(project_id)
        out_final_dir.mkdir(parents=True, exist_ok=True)
        out_final_thumb_dir = final_thumb_dir(project_id)
        out_final_thumb_dir.mkdir(parents=True, exist_ok=True)

        mask_image = template.main_black_mask_path if any(w in template.unique_name for w in ("mug", "cup")) else template.main_white_mask_path

        # Generate new final image
        try:
            new_final_path_str = await run_in_threadpool(
                generate_final_image_from_mask,
                template.base_image_path,
                template.main_white_mask_path,
                str(overlay_temp_abs),
                str(out_final_dir),
                template.unique_name
            )
        except Exception:
            safe_delete(overlay_temp_abs)
            logger.exception("Failed to regenerate final image for product_image_id=%s", product_image_id)
            raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to regenerate final image")

        new_final_abs = Path(new_final_path_str)

        color_temp = None
        try:
            if template.color_preview_path and template.color_white_mask_path and color:
                color_temp = await run_in_threadpool(
                    colorize_preview_with_mask,
                    abs_from_db(template.color_preview_path) or Path(template.color_preview_path),
                    abs_from_db(template.color_white_mask_path) or Path(template.color_white_mask_path),
                    color,
                    out_final_dir,
                )
                await run_in_threadpool(
                    apply_color_layer_to_image,
                    new_final_abs,
                    Path(color_temp),
                )
                safe_delete(Path(color_temp))
                color_temp = None
        except Exception:
            safe_delete(new_final_abs)
            safe_delete(Path(color_temp) if color_temp else None)
            safe_delete(overlay_temp_abs)
            logger.exception("Failed to apply color layer during regeneration")
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="Failed to apply color layer",
            )

        # Create new thumbnail
        try:
            new_thumb_abs = create_thumbnail_from_path(new_final_abs, out_final_thumb_dir)
        except Exception:
            safe_delete(new_final_abs)
            safe_delete(overlay_temp_abs)
            logger.exception("Failed to create new thumbnail for %s", new_final_abs)
            raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to regenerate final thumbnail")

        # Optional resized product image
        new_product_abs = None
        width = template.width
        height = template.height
        if width and height:
            try:
                w = int(width)
                h = int(height)
                new_product_abs = resize_image_to_size(new_final_abs, out_final_dir, w, h)
            except Exception:
                logger.exception("Failed to create resized product image")

        # Resolve old files to absolute paths for cleanup after DB commit
        def _abs_from_db_path(p: Optional[str]) -> Optional[Path]:
            return abs_from_db(p) if p else None

        old_final_abs = _abs_from_db_path(img_row.final_image_path)
        old_thumb_abs = _abs_from_db_path(img_row.final_image_thumbnail_path)
        old_prod_abs = _abs_from_db_path(img_row.final_image_product_path)

        # Update DB atomically
        def _rel(p: Optional[Path]) -> Optional[str]:
            return rel_path(p) if p else None

        try:
            img_row.final_image_path = _rel(new_final_abs)
            img_row.final_image_thumbnail_path = _rel(new_thumb_abs)
            img_row.final_image_product_path = _rel(new_product_abs) if new_product_abs else None

            db.add(img_row)
            db.commit()
            db.refresh(img_row)
        except Exception:
            db.rollback()
            # Cleanup newly created files
            safe_delete(new_final_abs)
            safe_delete(new_thumb_abs)
            safe_delete(new_product_abs)
            safe_delete(overlay_temp_abs)
            logger.exception("Failed to update DB while regenerating product image id=%s", product_image_id)
            raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to update database")

        # Delete old files after successful commit (best-effort)
        safe_delete(old_final_abs)
        safe_delete(old_thumb_abs)
        safe_delete(old_prod_abs)

        # Remove temporary overlay saved earlier
        safe_delete(overlay_temp_abs)

        return {
            "success": True,
            "message": "Final image regenerated successfully",
            "product_image_id": img_row.id,
            "project_master_id": img_row.project_master_id,
            "final_image_path": img_row.final_image_path,
            "final_image_thumbnail_path": img_row.final_image_thumbnail_path,
            "final_image_product_path": img_row.final_image_product_path,
        }