ubuntu-post-install/paintplus/backend/app/routers/edits.py

from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
from sqlalchemy.orm import Session
import json

from app.database import get_db
from app.models.project import Project
from app.models.edit import Edit
from app.schemas import EditRequest, EditResponse, StatusResponse
from app.services.edit_service import EditService
from app.config import settings

router = APIRouter(prefix="/edits", tags=["edits"])


async def process_edit_background(
    edit_id: int,
    project_id: int,
    request: EditRequest,
    db: Session
):
    """Background task to process edit"""
    edit_service = EditService()

    try:
        # Process the edit
        result_path = await edit_service.process_edit(
            project_id=project_id,
            edit_id=edit_id,
            prompt=request.prompt,
            mode=request.mode,
            selection_type=request.selection_type,
            bbox=request.bbox,
            feather_px=request.feather_px,
            selection_data=request.selection_data
        )

        # Update edit status
        edit = db.query(Edit).filter(Edit.id == edit_id).first()
        if edit:
            edit.status = "completed"
            db.commit()

    except Exception as e:
        # Update edit with error
        edit = db.query(Edit).filter(Edit.id == edit_id).first()
        if edit:
            edit.status = "failed"
            edit.error_message = str(e)
            db.commit()


@router.post("/projects/{project_id}/fix", response_model=EditResponse)
async def create_edit(
    project_id: int,
    request: EditRequest,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db)
):
    """
    Create a new edit request (Fix button)

    This endpoint accepts the selection data and prompt,
    then processes the edit in the background.
    """
    # Verify project exists
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Validate mode
    if request.mode not in ["A", "B"]:
        raise HTTPException(status_code=400, detail="Mode must be 'A' or 'B'")

    # Validate selection type
    if request.selection_type not in ["rectangle", "ellipse", "lasso"]:
        raise HTTPException(status_code=400, detail="Invalid selection type")

    # Create edit record
    edit = Edit(
        project_id=project_id,
        mode=request.mode,
        prompt=request.prompt,
        selection_type=request.selection_type,
        bbox_json=json.dumps(request.bbox),
        feather_px=request.feather_px,
        ai_provider=settings.ai_provider,
        status="pending"
    )
    db.add(edit)
    db.commit()
    db.refresh(edit)

    # Process edit in background
    background_tasks.add_task(
        process_edit_background,
        edit.id,
        project_id,
        request,
        db
    )

    return edit


@router.get("/{edit_id}", response_model=EditResponse)
def get_edit(
    edit_id: int,
    db: Session = Depends(get_db)
):
    """Get edit details and status"""
    edit = db.query(Edit).filter(Edit.id == edit_id).first()
    if not edit:
        raise HTTPException(status_code=404, detail="Edit not found")
    return edit


@router.post("/projects/{project_id}/revert/{edit_id}", response_model=StatusResponse)
def revert_to_edit(
    project_id: int,
    edit_id: int,
    db: Session = Depends(get_db)
):
    """Revert project to a specific edit"""
    # Verify project exists
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Verify edit exists and belongs to project
    edit = db.query(Edit).filter(
        Edit.id == edit_id,
        Edit.project_id == project_id
    ).first()
    if not edit:
        raise HTTPException(status_code=404, detail="Edit not found")

    # Revert
    edit_service = EditService()
    try:
        result_path = edit_service.revert_to_edit(project_id, edit_id)
        return StatusResponse(
            status="success",
            message=f"Reverted to edit {edit_id}",
            data={"image_url": f"/projects/{project_id}/current"}
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))


@router.post("/projects/{project_id}/reset", response_model=StatusResponse)
def reset_to_original(
    project_id: int,
    db: Session = Depends(get_db)
):
    """Reset project to original image"""
    # Verify project exists
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Reset
    edit_service = EditService()
    try:
        result_path = edit_service.reset_to_original(project_id)
        return StatusResponse(
            status="success",
            message="Reset to original image",
            data={"image_url": f"/projects/{project_id}/current"}
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))