import os
from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
import boto3

# ====== CONFIG FROM ENV ======
DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET", "milestone")
DO_ACCESS_KEY = os.getenv("DO_SPACES_ACCESS_KEY")
DO_SECRET_KEY = os.getenv("DO_SPACES_SECRET_KEY")

DO_SPACES_ENDPOINT = f"https://{DO_SPACES_REGION}.digitaloceanspaces.com"

# ====== INIT FASTAPI ======
app = FastAPI()

# Allow CORS for testing/frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def get_spaces_client():
    """Return a boto3 client for DigitalOcean Spaces"""
    return boto3.client(
        "s3",
        region_name=DO_SPACES_REGION,
        endpoint_url=DO_SPACES_ENDPOINT,
        aws_access_key_id=DO_ACCESS_KEY,
        aws_secret_access_key=DO_SECRET_KEY
    )


@app.get("/list-images")
def list_images(category: str = Query(
    None,
    description="Category name to list images for. If not provided, returns all available categories and their items."
)):
    client = get_spaces_client()
    base_prefix = "faceswap/target/"

    def get_items_for_category(cat):
        original_prefix = f"{base_prefix}{cat}/original/"
        thumb_prefix = f"{base_prefix}{cat}/thumb/"

        original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
        thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)

        original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
        thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]

        items = []
        for idx, filename in enumerate(sorted(original_files), start=1):
            if filename in thumb_files:
                items.append({
                    "categoryId": f"{cat.lower()}image_{idx}",
                    "url": f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/{original_prefix}{filename}",
                    "thumburl": f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/{thumb_prefix}{filename}"
                })
        return items

    if not category:
        resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
        categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
        all_data = []
        for cat in categories:
            items = get_items_for_category(cat)
            all_data.append({
                "categoryName": cat,
                "items": items
            })
        return {"categories": all_data, "note": "All categories and their items. Use /list-images?category=<category> to get a specific category."}

    items = get_items_for_category(category)
    return {
        "categoryName": category,
        "items": items,
        "note": "The URLs returned are public and do not require a bearer token."
    }


@app.get("/")
def home():
    return {"message": "FastAPI app running on Render with DigitalOcean Spaces"}