from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import requests
import os
import uuid
import shutil
from tempfile import NamedTemporaryFile

app = FastAPI()

# Mount a static directory to serve files
STATIC_DIR = "/app/static"
os.makedirs(STATIC_DIR, exist_ok=True)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

# API endpoint for ChangeClothesAI
CHANGE_CLOTHES_API = os.environ.get('URL')

@app.post("/try-on/")
async def try_on(
    human_img: UploadFile = File(...),
    garment_img: UploadFile = File(...),
    garment_desc: str = Form("dresses"),
    category: str = Form("dresses")
):
    try:
        # Generate unique filenames
        human_filename = f"{uuid.uuid4()}_{human_img.filename}"
        garment_filename = f"{uuid.uuid4()}_{garment_img.filename}"

        # Save files temporarily in the static directory
        human_path = os.path.join(STATIC_DIR, human_filename)
        garment_path = os.path.join(STATIC_DIR, garment_filename)

        try:
            # Save human image
            with open(human_path, "wb") as f:
                shutil.copyfileobj(human_img.file, f)
            # Save garment image
            with open(garment_path, "wb") as f:
                shutil.copyfileobj(garment_img.file, f)

            # Generate public URLs (assuming the Space is hosted at a public URL)
            # Replace YOUR_SPACE_URL with the actual Hugging Face Space URL
            base_url = "https://tejani-tryapi.hf.space"  # e.g., https://your-username-your-space.hf.space
            human_url = f"{base_url}/static/{human_filename}"
            garment_url = f"{base_url}/static/{garment_filename}"
            
            # Call the ChangeClothesAI API
            response = requests.post(CHANGE_CLOTHES_API, headers=headers, data=data)

            # Check if the API call was successful
            if response.status_code != 200:
                raise HTTPException(status_code=response.status_code, detail="Error calling ChangeClothesAI API")

            # Delete temporary files
            os.remove(human_path)
            os.remove(garment_path)

            return response.json()

        except Exception as e:
            # Clean up in case of any errors
            if os.path.exists(human_path):
                os.remove(human_path)
            if os.path.exists(garment_path):
                os.remove(garment_path)
            raise HTTPException(status_code=500, detail=str(e))

        finally:
            human_img.file.close()
            garment_img.file.close()

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing files: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)