main.py

import os
import time
import uuid
import shutil
from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from datetime import timezone, timedelta

from config import UPLOAD_DIR, PROJECT_NAME
from database import Base, engine, get_db
from models import Project, User
from schemas import PipelineResponse, ProjectResponse
from crud import (
    create_project,
    update_project_ready,
    update_project_error,
    get_all_projects,
    get_project_by_file_id,
    save_project_state,
    get_saved_state,
    delete_project,
)
from pipeline import run_pipeline
from auth_utils import get_current_user
from auth import router as auth_router

Base.metadata.create_all(bind=engine)

app = FastAPI(title=PROJECT_NAME)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(auth_router)


@app.get("/")
def root():
    return {"message": f"{PROJECT_NAME} API is running"}


@app.get("/health")
def health_check():
    return {"status": "ok", "project": PROJECT_NAME}


@app.post("/upload")
async def upload_image(
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    allowed_types = ["image/jpeg", "image/png", "image/webp"]
    if file.content_type not in allowed_types:
        raise HTTPException(400, "Only JPEG, PNG, WEBP images allowed")

    file_id   = str(uuid.uuid4())
    extension = file.filename.split(".")[-1]
    save_path = os.path.join(UPLOAD_DIR, f"{file_id}.{extension}")

    with open(save_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    project = create_project(db, file.filename, file_id, save_path, current_user.id)

    start = time.time()
    try:
        result = run_pipeline(save_path)
        processing_ms = int((time.time() - start) * 1000)
        update_project_ready(db, file_id, len(result["layers"]), processing_ms)
    except Exception as e:
        update_project_error(db, file_id)
        raise HTTPException(500, f"Pipeline failed: {str(e)}")

    return JSONResponse({
        "project_id":        project.id,
        "file_id":           file_id,
        "filename":          file.filename,
        "image_w":           result["image_w"],
        "image_h":           result["image_h"],
        "background_base64": result["background_base64"],
        "original_base64":   result["original_base64"],
        "layers":            result["layers"],
        "processing_time_s": result["processing_time_s"],
    })


@app.get("/projects")
def list_projects(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    projects = get_all_projects(db, current_user.id)
    return [
        {
            "id":            p.id,
            "filename":      p.filename,
            "file_id":       p.file_id,
            "layer_count":   p.layer_count,
            "status":        p.status,
            "processing_ms": p.processing_ms,
            "created_at": (
                p.created_at.replace(tzinfo=timezone.utc)
                .astimezone(timezone(timedelta(hours=5, minutes=30)))
                .isoformat()
            ),
            "thumbnail_b64": p.thumbnail_b64,
        }
        for p in projects
    ]


@app.get("/projects/{file_id}")
def get_project(
    file_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    project = get_project_by_file_id(db, file_id, current_user.id)
    if not project:
        raise HTTPException(404, "Project not found")
    return {
        "id":            project.id,
        "filename":      project.filename,
        "file_id":       project.file_id,
        "layer_count":   project.layer_count,
        "status":        project.status,
        "processing_ms": project.processing_ms,
        "created_at": (
            project.created_at.replace(tzinfo=timezone.utc)
            .astimezone(timezone(timedelta(hours=5, minutes=30)))
            .isoformat()
        ),
    }


@app.post("/inpaint")
async def inpaint_element(
    request: dict,
    current_user: User = Depends(get_current_user),
):
    import base64
    from inpaint import inpaint_region, unload_inpaint_model, create_mask_from_bbox

    file_id = request.get("file_id")
    x, y    = request.get("x", 0), request.get("y", 0)
    w, h    = request.get("w", 0), request.get("h", 0)

    if not file_id:
        raise HTTPException(400, "file_id is required")

    upload_files = os.listdir(UPLOAD_DIR)
    image_path   = None
    for f in upload_files:
        if f.startswith(file_id):
            image_path = os.path.join(UPLOAD_DIR, f)
            break

    if not image_path:
        raise HTTPException(404, "Original image not found")

    try:
        mask   = create_mask_from_bbox(image_path, x, y, w, h)
        result = inpaint_region(image_path, mask)
        unload_inpaint_model()

        out_path = os.path.join("outputs", f"inpaint_{file_id}.jpg")
        result.save(out_path, quality=95)

        with open(out_path, "rb") as f:
            b64 = base64.b64encode(f.read()).decode("utf-8")

        return JSONResponse({"inpainted_image_base64": b64})

    except Exception as e:
        unload_inpaint_model()
        raise HTTPException(500, f"Inpainting failed: {str(e)}")


@app.get("/projects/{file_id}/layers")
async def get_project_layers(
    file_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    project = get_project_by_file_id(db, file_id, current_user.id)
    if not project:
        raise HTTPException(404, "Project not found")

    if not os.path.exists(project.upload_path):
        raise HTTPException(404, "Original file no longer exists")

    try:
        result = run_pipeline(project.upload_path)
        return JSONResponse({
            "project_id":        project.id,
            "file_id":           file_id,
            "filename":          project.filename,
            "image_w":           result["image_w"],
            "image_h":           result["image_h"],
            "background_base64": result["background_base64"],
            "original_base64":   result["original_base64"],
            "layers":            result["layers"],
            "processing_time_s": result["processing_time_s"],
        })
    except Exception as e:
        raise HTTPException(500, f"Pipeline failed: {str(e)}")


@app.post("/projects/{file_id}/save")
async def save_project(
    file_id: str,
    request: dict,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    project = get_project_by_file_id(db, file_id, current_user.id)
    if not project:
        raise HTTPException(404, "Project not found")

    saved_state   = request.get("saved_state")
    thumbnail_b64 = request.get("thumbnail_b64")
    layer_count   = request.get("layer_count", project.layer_count)

    if not saved_state:
        raise HTTPException(400, "saved_state is required")

    save_project_state(db, file_id, saved_state, thumbnail_b64, layer_count, current_user.id)
    return JSONResponse({"status": "saved", "file_id": file_id})


@app.get("/projects/{file_id}/saved")
async def load_saved_project(
    file_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    data = get_saved_state(db, file_id, current_user.id)
    if not data:
        raise HTTPException(404, "Project not found")
    return JSONResponse(data)


@app.delete("/projects/{file_id}")
async def remove_project(
    file_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    project = get_project_by_file_id(db, file_id, current_user.id)
    if not project:
        raise HTTPException(404, "Project not found")

    if project.upload_path and os.path.exists(project.upload_path):
        try:
            os.remove(project.upload_path)
        except:
            pass

    inpaint_path = os.path.join("outputs", f"inpaint_{file_id}.jpg")
    if os.path.exists(inpaint_path):
        try:
            os.remove(inpaint_path)
        except:
            pass

    delete_project(db, file_id, current_user.id)
    return {"status": "deleted"}