from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
import requests
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import os
import uuid
from pathlib import Path
import logging

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

app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static directory for serving files
STATIC_DIR = Path("/home/user/app/static")
STATIC_DIR.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

# Define the request model
class TryOnRequest(BaseModel):
    garmentDesc: str
    category: str

# Allowed file extensions
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png"}

# Helper function to save file and generate public URL
async def save_file_and_get_url(file: UploadFile) -> str:
    try:
        # Validate file extension
        file_extension = f".{file.filename.split('.')[-1].lower()}"
        if file_extension not in ALLOWED_EXTENSIONS:
            logger.error(f"Invalid file extension for {file.filename}: {file_extension}")
            raise HTTPException(status_code=400, detail=f"File extension {file_extension} not allowed. Use JPG or PNG.")

        # Generate unique filename
        unique_filename = f"{uuid.uuid4()}{file_extension}"
        file_path = STATIC_DIR / unique_filename
        
        # Save file
        logger.info(f"Saving file to {file_path}")
        with file_path.open("wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # Verify file exists
        if not file_path.exists():
            logger.error(f"File {file_path} was not saved correctly")
            raise HTTPException(status_code=500, detail="Failed to save file")
        
        public_url = f"https://tejani-tryapi.hf.space/static/{unique_filename}"
        logger.info(f"Generated public URL: {public_url}")
        
        # Test URL accessibility
        try:
            response = requests.head(public_url, timeout=5)
            if response.status_code != 200:
                logger.warning(f"Public URL {public_url} returned status {response.status_code}")
            else:
                logger.info(f"Public URL {public_url} is accessible")
        except requests.exceptions.RequestException as e:
            logger.error(f"Failed to access public URL {public_url}: {str(e)}")
        
        return public_url
    except Exception as e:
        logger.error(f"Error in save_file_and_get_url: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")

# Endpoint to handle file uploads and proxy the request
@app.post("/try-on")
async def try_on(
    human_img: UploadFile = File(...),
    garment: UploadFile = File(...),
    garment_desc: str = "",
    category: str = "upper_body"
):
    try:
        # Save files and get public URLs
        human_img_url = await save_file_and_get_url(human_img)
        garment_url = await save_file_and_get_url(garment)
        
        # Original API endpoint
        url = "https://changeclothesai.online/api/try-on/edge"
        
        headers = {
            "accept": "*/*",
            "f": "sdasdwsKaVgUoxa5j1jzcFtziPx",
        }
        
        data = {
            "humanImg": human_img_url,
            "garment": garment_url,
            "garmentDesc": garment_desc,
            "category": category
        }
        
        logger.info(f"Forwarding request to {url} with data: {data}")
        
        # Forward request to the original API
        response = requests.post(url, headers=headers, cookies={}, data=data)
        response.raise_for_status()
        
        return {
            "status_code": response.status_code,
            "response": response.json() if response.headers.get('content-type') == 'application/json' else response.text,
            "human_img_url": human_img_url,
            "garment_url": garment_url
        }
    except requests.exceptions.RequestException as e:
        logger.error(f"Error forwarding request: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error forwarding request: {str(e)}")
    except Exception as e:
        logger.error(f"Error in try_on endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")

# Health check endpoint
@app.get("/")
async def root():
    return {"message": "FastAPI proxy for try-on API with file upload is running"}

# Debug endpoint to list stored files
@app.get("/list-files")
async def list_files():
    try:
        files = [str(f) for f in STATIC_DIR.glob("*") if f.is_file()]
        logger.info(f"Files in {STATIC_DIR}: {files}")
        return {"files": files}
    except Exception as e:
        logger.error(f"Error listing files: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error listing files: {str(e)}")