PaintPlus/backend/app/services/edit_service.py

import os
import json
from pathlib import Path
from typing import Dict, Optional
from datetime import datetime
from PIL import Image

from app.models.edit import Edit
from app.models.project import Project
from app.services.ai_provider import get_ai_provider
from app.utils.image_processing import (
    bytes_to_image,
    image_to_bytes,
    crop_patch,
    blend_patch,
    insert_patch,
    create_mask_from_selection,
    resize_for_ai,
    scale_bbox
)
from app.config import settings


class EditService:
    """Service for handling image edits"""

    def __init__(self, data_dir: str = None):
        self.data_dir = data_dir or settings.data_dir
        self.ai_provider = get_ai_provider()

    def get_project_dir(self, project_id: int) -> Path:
        """Get project directory path"""
        return Path(self.data_dir) / "projects" / str(project_id)

    def get_edit_dir(self, project_id: int, edit_id: int) -> Path:
        """Get edit history directory path"""
        return self.get_project_dir(project_id) / "history" / str(edit_id)

    def ensure_project_dir(self, project_id: int):
        """Ensure project directory structure exists"""
        project_dir = self.get_project_dir(project_id)
        project_dir.mkdir(parents=True, exist_ok=True)
        (project_dir / "history").mkdir(exist_ok=True)

    def get_current_image_path(self, project_id: int) -> Path:
        """Get path to current image"""
        return self.get_project_dir(project_id) / "current.png"

    def get_original_image_path(self, project_id: int) -> Path:
        """Get path to original image"""
        return self.get_project_dir(project_id) / "original.png"

    async def process_edit(
        self,
        project_id: int,
        edit_id: int,
        prompt: str,
        mode: str,
        selection_type: str,
        bbox: Dict[str, int],
        feather_px: int,
        selection_data: Optional[Dict] = None
    ) -> str:
        """
        Process an edit request

        Args:
            project_id: Project ID
            edit_id: Edit ID
            prompt: AI prompt
            mode: "A" or "B"
            selection_type: "rectangle", "ellipse", or "lasso"
            bbox: Bounding box {x, y, width, height}
            feather_px: Feather radius in pixels
            selection_data: Additional selection data (for lasso)

        Returns:
            Path to the result image
        """
        # Create edit directory
        edit_dir = self.get_edit_dir(project_id, edit_id)
        edit_dir.mkdir(parents=True, exist_ok=True)

        # Load current image
        current_image_path = self.get_current_image_path(project_id)
        full_image = Image.open(current_image_path).convert('RGBA')

        # Crop patch from current image
        original_patch = crop_patch(full_image, bbox)

        # Save original patch
        original_patch.save(edit_dir / "patch_in.png")

        # Create mask based on selection type
        mask = create_mask_from_selection(
            bbox['width'],
            bbox['height'],
            selection_type,
            selection_data or {}
        )

        # Save mask
        mask.save(edit_dir / "mask.png")

        # Resize patch and mask for AI if needed
        patch_for_ai, scale = resize_for_ai(original_patch)
        mask_for_ai = mask.resize(patch_for_ai.size, Image.Resampling.LANCZOS)

        # Prepare full image for mode B
        full_image_bytes = None
        if mode == "B":
            full_image_for_ai, _ = resize_for_ai(full_image)
            full_image_bytes = image_to_bytes(full_image_for_ai)

        # Call AI provider
        regenerated_patch_bytes = await self.ai_provider.edit_image(
            patch_image_bytes=image_to_bytes(patch_for_ai),
            mask_image_bytes=image_to_bytes(mask_for_ai),
            prompt=prompt,
            mode=mode,
            full_image_bytes=full_image_bytes
        )

        # Convert regenerated patch back to PIL Image
        regenerated_patch = bytes_to_image(regenerated_patch_bytes)

        # Resize back to original patch size if scaled
        if scale != 1.0:
            regenerated_patch = regenerated_patch.resize(
                original_patch.size,
                Image.Resampling.LANCZOS
            )

        # Save regenerated patch
        regenerated_patch.save(edit_dir / "patch_out.png")

        # Blend regenerated patch with original using mask
        blended_patch = blend_patch(
            original_patch,
            regenerated_patch,
            mask,
            feather_px
        )

        # Insert blended patch back into full image
        result_image = insert_patch(full_image, blended_patch, bbox)

        # Save result
        result_path = edit_dir / "result.png"
        result_image.save(result_path)

        # Update current image
        result_image.save(current_image_path)

        # Save metadata
        metadata = {
            'edit_id': edit_id,
            'project_id': project_id,
            'prompt': prompt,
            'mode': mode,
            'selection_type': selection_type,
            'bbox': bbox,
            'feather_px': feather_px,
            'selection_data': selection_data,
            'timestamp': datetime.utcnow().isoformat(),
            'ai_provider': settings.ai_provider
        }

        with open(edit_dir / "meta.json", 'w') as f:
            json.dump(metadata, f, indent=2)

        return str(result_path)

    def revert_to_edit(self, project_id: int, edit_id: int) -> str:
        """
        Revert project to a specific edit

        Args:
            project_id: Project ID
            edit_id: Edit ID to revert to

        Returns:
            Path to the reverted image
        """
        edit_dir = self.get_edit_dir(project_id, edit_id)
        result_path = edit_dir / "result.png"

        if not result_path.exists():
            raise FileNotFoundError(f"Edit {edit_id} result not found")

        # Copy result to current (preserve alpha channel)
        current_path = self.get_current_image_path(project_id)
        img = Image.open(result_path)
        # Preserve original mode to maintain transparency
        img.save(current_path, format='PNG')

        return str(current_path)

    def reset_to_original(self, project_id: int) -> str:
        """
        Reset project to original image

        Args:
            project_id: Project ID

        Returns:
            Path to the original image
        """
        original_path = self.get_original_image_path(project_id)
        current_path = self.get_current_image_path(project_id)

        if not original_path.exists():
            raise FileNotFoundError(f"Original image for project {project_id} not found")

        # Copy original to current (preserve alpha channel)
        img = Image.open(original_path)
        # Preserve original mode to maintain transparency
        img.save(current_path, format='PNG')

        return str(current_path)