from fastapi import FastAPI, Request, Depends, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os
from pathlib import Path

from seed_plans import seed_plans

from router.image import background_remove
from router.voice import clone
from router.try_on import try_clothes
from router.lip_sync import lip_sync
import torch
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

from starlette.middleware.base import BaseHTTPMiddleware
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title="AI Image Processing API",
    description="API for image processing with AI models",
    version="1.0.0",
    debug=True,
)

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
import traceback


class CustomCORSMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.method == "OPTIONS":
            response = Response()
        else:
            try:
                response = await call_next(request)
            except Exception as e:
                response = Response(
                    content=f"Internal Server Error: {str(e)}", status_code=500
                )

        # Get the origin from the request
        origin = request.headers.get("origin")

        # Allow specific origins (not wildcard) when using credentials
        allowed_origins = [
            "http://localhost:3000",
            "http://localhost:3001",
            "http://localhost:5173",
            "http://localhost:8080",
        ]

        if origin in allowed_origins:
            response.headers["Access-Control-Allow-Origin"] = (
                origin  # Specific origin, not "*"
            )

        response.headers["Access-Control-Allow-Credentials"] = (
            "true"  # Must be string "true"
        )
        response.headers["Access-Control-Allow-Methods"] = "*"
        response.headers["Access-Control-Allow-Headers"] = "*"
        response.headers["Access-Control-Max-Age"] = "86400"

        return response


app.add_middleware(CustomCORSMiddleware)

# Add CORS middleware for Next.js frontend
app.add_middleware(
    CORSMiddleware,
    # allow_origins=["http://localhost:3000"],  # NextJS default port
    # allow_credentials=True,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=True,  # Must be False when using "*"
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables
birefnet = None
birefnet_lite = None
transform_image = None


# Include routers
# app.include_router(ai_image_router)
app.include_router(background_remove.router)
app.include_router(clone.router)
app.include_router(try_clothes.router)

# Create temp directory for file storage if it doesn't exist
temp_dir = Path("./temp")
temp_dir.mkdir(exist_ok=True)


@app.on_event("startup")
def startup_event():
    """Create database tables if they don't exist."""
    print("DB URL:", os.getenv("DATABASE_URL"))
    try_clothes.initialize_models()
    lip_sync.load_model("wav2lip_checkpoints/wav2lip_gan.pth")

    logger.info("Loading BiRefNet models...")
    torch.set_float32_matmul_precision("medium")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load both models and store in app state
    app.state.birefnet = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet", trust_remote_code=True
    )
    app.state.birefnet.to(device)

    app.state.birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet_lite", trust_remote_code=True
    )
    app.state.birefnet_lite.to(device)

    app.state.transform_image = transforms.Compose(
        [
            transforms.Resize((768, 768)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )

    logger.info("Models loaded successfully")


@app.get("/")
async def root():
    return {
        "message": "Voice Processing API",
        "version": "1.0.0",
        "endpoints": {
            "voice_cloning": "/voice/clone",
            "speech_enhancement": "/voice/speech-enhancement",
            "speech_separation": "/voice/speech-separation",
            "target_speaker_extraction": "/voice/target-speaker-extraction",
            "speech_super_resolution": "/voice/speech-super-resolution",
            "health_check": "/voice/health",
            "cleanup": "/voice/cleanup",
            "download": "/voice/download/{filename}",
        },
    }


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


if __name__ == "__main__":
    # Start server with hot reload during development
    # uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        reload_dirs=["./ai-api"],
        log_level="debug",  # Add this line for debug logs
    )