app.py

import os
import io
import json
import traceback
from datetime import datetime,timedelta
from typing import Optional
import time
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pymongo import MongoClient
import gridfs
from bson.objectid import ObjectId
from PIL import Image
from fastapi.concurrency import run_in_threadpool
import shutil
import firebase_admin
from firebase_admin import credentials, auth
from PIL import Image
from huggingface_hub import InferenceClient

# ---------------------------------------------------------------------
# Load Firebase Config from env (stringified JSON)
# ---------------------------------------------------------------------
firebase_config_json = os.getenv("firebase_config")
if not firebase_config_json:
    raise RuntimeError("❌ Missing Firebase config in environment variable 'firebase_config'")

try:
    firebase_creds_dict = json.loads(firebase_config_json)
    cred = credentials.Certificate(firebase_creds_dict)
    firebase_admin.initialize_app(cred)
except Exception as e:
    raise RuntimeError(f"Failed to initialize Firebase Admin SDK: {e}")

# ---------------------------------------------------------------------
# Hugging Face setup
# ---------------------------------------------------------------------
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN not set in environment variables")

hf_client = InferenceClient(token=HF_TOKEN)

# ---------------------------------------------------------------------
# MongoDB setup
# ---------------------------------------------------------------------
MONGODB_URI=os.getenv("MONGODB_URI")
DB_NAME = "polaroid_db"

mongo = MongoClient(MONGODB_URI)
db = mongo[DB_NAME]
fs = gridfs.GridFS(db)
logs_collection = db["logs"]

# ---------------------------------------------------------------------
# FastAPI app setup
# ---------------------------------------------------------------------
app = FastAPI(title="Qwen Image Edit API with Firebase Auth")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------
# Auth dependency
# ---------------------------------------------------------------------
async def verify_firebase_token(request: Request):
    """Middleware-like dependency to verify Firebase JWT from Authorization header."""
    auth_header = request.headers.get("Authorization")
    if not auth_header or not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")

    id_token = auth_header.split("Bearer ")[1]
    try:
        decoded_token = auth.verify_id_token(id_token)
        request.state.user = decoded_token
        return decoded_token
    except Exception as e:
        raise HTTPException(status_code=401, detail=f"Invalid or expired Firebase token: {e}")

# ---------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------
class HealthResponse(BaseModel):
    status: str
    db: str
    model: str

# --------------------- UTILS ---------------------
def resize_image_if_needed(img: Image.Image, max_size=(1024, 1024)) -> Image.Image:
    """
    Resize image to fit within max_size while keeping aspect ratio.
    """
    if img.width > max_size[0] or img.height > max_size[1]:
        img.thumbnail(max_size, Image.ANTIALIAS)
    return img

def prepare_image(file_bytes: bytes) -> Image.Image:
    """
    Open image and resize if larger than 1024x1024
    """
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")

    # ✅ MIN SIZE CHECK
    if img.width < 200 or img.height < 200:
        raise HTTPException(
            status_code=400,
            detail="Image size is below 200x200 pixels. Please upload a larger image."
        )
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
    img = resize_image_if_needed(img, max_size=(1024, 1024))
    return img
    
MAX_COMPRESSED_SIZE = 2 * 1024 * 1024  # 2 MB
def compress_pil_image_to_2mb(
    pil_img: Image.Image,
    max_dim: int = 1280
) -> bytes:
    """
    Resize + compress PIL image to <= 2MB.
    Returns JPEG bytes.
    """
    img = pil_img.convert("RGB")

    # Resize (maintain aspect ratio)
    img.thumbnail((max_dim, max_dim), Image.LANCZOS)

    quality = 85
    buffer = io.BytesIO()

    while quality >= 40:
        buffer.seek(0)
        buffer.truncate()

        img.save(
            buffer,
            format="JPEG",
            quality=quality,
            optimize=True,
            progressive=True
        )

        if buffer.tell() <= MAX_COMPRESSED_SIZE:
            break

        quality -= 5

    return buffer.getvalue()
# ---------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------
@app.get("/health", response_model=HealthResponse)
def health():
    """Public health check"""
    mongo.admin.command("ping")
    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")

@app.post("/generate")
async def generate(
    prompt: str = Form(...),
    image1: UploadFile = File(...),
    image2: Optional[UploadFile] = File(None),
    user_id: Optional[str] = Form(None),
    category_id: Optional[str] = Form(None),
    user=Depends(verify_firebase_token)
):
    start_time = time.time()

    # -------------------------
    # 1. VALIDATE & READ IMAGES
    # -------------------------
    try:
        img1_bytes = await image1.read()
        pil_img1 = prepare_image(img1_bytes)
        input1_id = fs.put(
            img1_bytes,
            filename=image1.filename,
            contentType=image1.content_type,
            metadata={"role": "input"}
        )
    except Exception as e:
        raise HTTPException(400, f"Failed to read first image: {e}")

    img2_bytes = None
    input2_id = None
    pil_img2 = None

    if image2:
        try:
            img2_bytes = await image2.read()
            pil_img2 = prepare_image(img2_bytes)
            input2_id = fs.put(
                img2_bytes,
                filename=image2.filename,
                contentType=image2.content_type,
                metadata={"role": "input"}
            )
        except Exception as e:
            raise HTTPException(400, f"Failed to read second image: {e}")

    # -------------------------
    # 2. COMBINE IF NEEDED
    # -------------------------
    if pil_img2:
        total_width = pil_img1.width + pil_img2.width
        max_height = max(pil_img1.height, pil_img2.height)
        combined_img = Image.new("RGB", (total_width, max_height))
        combined_img.paste(pil_img1, (0, 0))
        combined_img.paste(pil_img2, (pil_img1.width, 0))
    else:
        combined_img = pil_img1

    # -------------------------
    # 3. CATEGORY CLICK LOGIC
    # -------------------------
    if user_id and category_id:
        try:
            admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
            admin_db = admin_client["adminPanel"]
    
            categories_col = admin_db.categories
            media_clicks_col = admin_db.media_clicks
    
            # Validate user_oid & category_oid
            user_oid = ObjectId(user_id)
            category_oid = ObjectId(category_id)
    
            # Check category exists
            category_doc = categories_col.find_one({"_id": category_oid})
            if not category_doc:
                raise HTTPException(400, f"Invalid category_id: {category_id}")
    
            now = datetime.utcnow()
    
            # Normalize dates (UTC midnight)
            today_date = datetime(now.year, now.month, now.day)
            yesterday_date = today_date - timedelta(days=1)
    
            # --------------------------------------------------
            # AI EDIT USAGE TRACKING (GLOBAL PER USER)
            # --------------------------------------------------
            media_clicks_col.update_one(
                {"userId": user_oid},
                {
                    "$setOnInsert": {
                        "createdAt": now,
                        "ai_edit_daily_count": []
                    },
                    "$set": {
                        "ai_edit_last_date": now,
                        "updatedAt": now
                    },
                    "$inc": {
                        "ai_edit_complete": 1
                    }
                },
                upsert=True
            )
    
            # --------------------------------------------------
            # DAILY COUNT LOGIC
            # --------------------------------------------------
            now = datetime.utcnow()
            today_date = datetime(now.year, now.month, now.day)
            
            doc = media_clicks_col.find_one(
                {"userId": user_oid},
                {"ai_edit_daily_count": 1}
            )
            
            daily_entries = doc.get("ai_edit_daily_count", []) if doc else []
            
            # Build UNIQUE date -> count map
            daily_map = {}
            for entry in daily_entries:
                d = entry["date"]
                d = datetime(d.year, d.month, d.day) if isinstance(d, datetime) else d
                daily_map[d] = entry["count"]  # overwrite = no duplicates
            
            # Find last known date
            last_date = max(daily_map.keys()) if daily_map else today_date
            
            # Fill ALL missing days with 0
            next_day = last_date + timedelta(days=1)
            while next_day < today_date:
                daily_map.setdefault(next_day, 0)
                next_day += timedelta(days=1)
            
            # Mark today as used (binary)
            daily_map[today_date] = 1
            
            # Rebuild list (OLD → NEW)
            final_daily_entries = [
                {"date": d, "count": daily_map[d]}
                for d in sorted(daily_map.keys())
            ]
            
            # Keep last 32 days only
            final_daily_entries = final_daily_entries[-32:]
            
            # ATOMIC REPLACE (NO PUSH)
            media_clicks_col.update_one(
                {"userId": user_oid},
                {
                    "$set": {
                        "ai_edit_daily_count": final_daily_entries,
                        "ai_edit_last_date": now,
                        "updatedAt": now
                    }
                }
            )


            # --------------------------------------------------
            # CATEGORY CLICK LOGIC
            # --------------------------------------------------
            update_res = media_clicks_col.update_one(
                {"userId": user_oid, "categories.categoryId": category_oid},
                {
                    "$set": {
                        "updatedAt": now,
                        "categories.$.lastClickedAt": now
                    },
                    "$inc": {
                        "categories.$.click_count": 1
                    }
                }
            )
    
            # If category does not exist → push new
            if update_res.matched_count == 0:
                media_clicks_col.update_one(
                    {"userId": user_oid},
                    {
                        "$set": {"updatedAt": now},
                        "$push": {
                            "categories": {
                                "categoryId": category_oid,
                                "click_count": 1,
                                "lastClickedAt": now
                            }
                        }
                    },
                    upsert=True
                )
    
        except Exception as e:
            print("CATEGORY_LOG_ERROR:", e)
    # -------------------------
    # 4. HF INFERENCE
    # -------------------------
    try:
        pil_output = hf_client.image_to_image(
            image=combined_img,
            prompt=prompt,
            model="Qwen/Qwen-Image-Edit"
        )
    except Exception as e:
        response_time_ms = round((time.time() - start_time) * 1000)
        logs_collection.insert_one({
            "timestamp": datetime.utcnow(),
            "status": "failure",
            "input1_id": str(input1_id),
            "input2_id": str(input2_id) if input2_id else None,
            "prompt": prompt,
            "user_email": user.get("email"),
            "error": str(e),
            "response_time_ms": response_time_ms
        })
        raise HTTPException(500, f"Inference failed: {e}")

    # -------------------------
    # 5. SAVE OUTPUT IMAGE
    # -------------------------
    out_buf = io.BytesIO()
    pil_output.save(out_buf, format="PNG")
    out_bytes = out_buf.getvalue()

    out_id = fs.put(
        out_bytes,
        filename=f"result_{input1_id}.png",
        contentType="image/png",
        metadata={
            "role": "output",
            "prompt": prompt,
            "input1_id": str(input1_id),
            "input2_id": str(input2_id) if input2_id else None,
            "user_email": user.get("email"),
        }
    )
    # -------------------------
    # 5b. SAVE COMPRESSED IMAGE
    # -------------------------
    compressed_bytes = compress_pil_image_to_2mb(
        pil_output,
        max_dim=1280
    )
    
    compressed_id = fs.put(
        compressed_bytes,
        filename=f"result_{input1_id}_compressed.jpg",
        contentType="image/jpeg",
        metadata={
            "role": "output_compressed",
            "original_output_id": str(out_id),
            "prompt": prompt,
            "user_email": user.get("email")
        }
    )


    response_time_ms = round((time.time() - start_time) * 1000)

    # -------------------------
    # 6. LOG SUCCESS
    # -------------------------
    logs_collection.insert_one({
        "timestamp": datetime.utcnow(),
        "status": "success",
        "input1_id": str(input1_id),
        "input2_id": str(input2_id) if input2_id else None,
        "output_id": str(out_id),
        "prompt": prompt,
        "user_email": user.get("email"),
        "response_time_ms": response_time_ms
    })

    return JSONResponse({
        "output_image_id": str(out_id),
        "user": user.get("email"),
        "response_time_ms": response_time_ms,
        "Compressed_Image_URL": (
            f"https://logicgoinfotechspaces-polaroidimage.hf.space/image/{compressed_id}"
        )
    })
####################---------------------------------------------------------------------------------------------------###
###-----OLD CODE--------------####
# @app.post("/generate")
# async def generate(
#     prompt: str = Form(...),
#     image1: UploadFile = File(...),
#     image2: Optional[UploadFile] = File(None),
#     user=Depends(verify_firebase_token)
# ):
#     start_time = time.time()

#     # PREPARE IMAGE(S)
#     try:
#         img1_bytes = await image1.read()
#         pil_img1 = prepare_image(img1_bytes)
#         input1_id = fs.put(
#             img1_bytes,
#             filename=image1.filename,
#             contentType=image1.content_type,
#             metadata={"role": "input"}
#         )
#     except Exception as e:
#         raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")

#     img2_bytes = None
#     input2_id = None
#     pil_img2 = None

#     if image2:
#         try:
#             img2_bytes = await image2.read()
#             pil_img2 = prepare_image(img2_bytes)
#             input2_id = fs.put(
#                 img2_bytes,
#                 filename=image2.filename,
#                 contentType=image2.content_type,
#                 metadata={"role": "input"}
#             )
#         except Exception as e:
#             raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")

#     # COMBINE IF NEEDED
#     if pil_img2:
#         total_width = pil_img1.width + pil_img2.width
#         max_height = max(pil_img1.height, pil_img2.height)
#         combined_img = Image.new("RGB", (total_width, max_height))
#         combined_img.paste(pil_img1, (0, 0))
#         combined_img.paste(pil_img2, (pil_img1.width, 0))
#     else:
#         combined_img = pil_img1

#     # INFERENCE
#     try:
#         pil_output = hf_client.image_to_image(
#             image=combined_img,
#             prompt=prompt,
#             model="Qwen/Qwen-Image-Edit"
#         )
#     except Exception as e:
#         # LOG FAILURE
#         response_time_ms = round((time.time() - start_time) * 1000)
#         logs_collection.insert_one({
#             "timestamp": datetime.utcnow(),
#             "status": "failure",
#             "input1_id": str(input1_id),
#             "input2_id": str(input2_id) if input2_id else None,
#             "prompt": prompt,
#             "user_email": user.get("email"),
#             "error": str(e),
#             "response_time_ms": response_time_ms
#         })
#         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")

#     # SAVE OUTPUT IMAGE
#     out_buf = io.BytesIO()
#     pil_output.save(out_buf, format="PNG")
#     out_bytes = out_buf.getvalue()

#     out_id = fs.put(
#         out_bytes,
#         filename=f"result_{input1_id}.png",
#         contentType="image/png",
#         metadata={
#             "role": "output",
#             "prompt": prompt,
#             "input1_id": str(input1_id),
#             "input2_id": str(input2_id) if input2_id else None,
#             "user_email": user.get("email"),
#         }
#     )

#     # FINAL RESPONSE TIME
#     response_time_ms = round((time.time() - start_time) * 1000)

#     # LOG SUCCESS - THIS WAS MISSING THE STATUS AND RESPONSE_TIME_MS FIELDS!
#     logs_collection.insert_one({
#         "timestamp": datetime.utcnow(),
#         "status": "success",  # ✅ FIXED: Added status field
#         "input1_id": str(input1_id),
#         "input2_id": str(input2_id) if input2_id else None,
#         "output_id": str(out_id),
#         "prompt": prompt,
#         "user_email": user.get("email"),  # ✅ FIXED: Changed from "user" to "user_email"
#         "response_time_ms": response_time_ms  # ✅ FIXED: Added response_time_ms field
#     })

#     return JSONResponse({
#         "output_image_id": str(out_id),
#         "user": user.get("email"),
#         "response_time_ms": response_time_ms
#     })



@app.get("/image/{image_id}")
def get_image(image_id: str, download: Optional[bool] = False):
    """Retrieve stored image by ID (no authentication required)."""
    try:
        oid = ObjectId(image_id)
        grid_out = fs.get(oid)
    except Exception:
        raise HTTPException(status_code=404, detail="Image not found")

    def iterfile():
        yield grid_out.read()

    headers = {}
    if download:
        headers["Content-Disposition"] = f'attachment; filename="{grid_out.filename}"'

    return StreamingResponse(
        iterfile(),
        media_type=grid_out.content_type or "application/octet-stream",
        headers=headers
    )


# ---------------------------------------------------------------------
# Run locally
# ---------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)