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)
# ---------------------------------------------------------------------
# MODEL SELECTION
# ---------------------------------------------------------------------
genai = None  # ✅ IMPORTANT: module-level declaration
MODEL = os.getenv("IMAGE_MODEL", "GEMINI").upper()
GEMINI_FORCE_CATEGORY_ID = "69368e741224bcb6bdb98076"
# ---------------------------------------------------------------------
# Gemini setup (ONLY used if MODEL == "GEMINI")
# ---------------------------------------------------------------------
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEMINI_IMAGE_MODEL = os.getenv("GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image")

# ---------------------------------------------------------------------
# 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
# ---------------------------------------------------------------------
# Lazy Gemini Initialization
# ---------------------------------------------------------------------
_genai_initialized = False
def init_gemini():
    global _genai_initialized, genai

    if _genai_initialized:
        return

    if not GEMINI_API_KEY:
        raise RuntimeError("❌ GEMINI_API_KEY not set")

    import google.generativeai as genai
    genai.configure(api_key=GEMINI_API_KEY)

    _genai_initialized = True

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()
    
def run_image_generation(
    image1: Image.Image,
    prompt: str,
    image2: Optional[Image.Image] = None,
    force_model: Optional[str] = None
) -> Image.Image:
    """
    Unified image generation interface.
    QWEN  -> merges images if image2 exists
    GEMINI -> passes images separately
    """

    effective_model = (force_model or MODEL).upper()

    # ---------------- QWEN ----------------
    if effective_model == "QWEN":

        # ✅ Merge images ONLY for QWEN
        if image2:
            total_width = image1.width + image2.width
            max_height = max(image1.height, image2.height)
            merged = Image.new("RGB", (total_width, max_height))
            merged.paste(image1, (0, 0))
            merged.paste(image2, (image1.width, 0))
        else:
            merged = image1

        return hf_client.image_to_image(
            image=merged,
            prompt=prompt,
            model="Qwen/Qwen-Image-Edit"
        )

    # ---------------- GEMINI ----------------
    elif effective_model == "GEMINI":
        init_gemini()
        model = genai.GenerativeModel(GEMINI_IMAGE_MODEL)

        parts = [prompt]

        def add_image(img: Image.Image):
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            buf.seek(0)
            parts.append({
                "mime_type": "image/png",
                "data": buf.getvalue()
            })

        add_image(image1)

        if image2:
            add_image(image2)

        response = model.generate_content(parts)

        # -------- SAFE IMAGE EXTRACTION --------
        import base64

        image_bytes = None
        for candidate in response.candidates:
            for part in candidate.content.parts:
                if hasattr(part, "inline_data") and part.inline_data:
                    data = part.inline_data.data
                    image_bytes = (
                        data if isinstance(data, (bytes, bytearray))
                        else base64.b64decode(data)
                    )
                    break
            if image_bytes:
                break

        if not image_bytes:
            raise RuntimeError("Gemini did not return an image")

        img = Image.open(io.BytesIO(image_bytes))
        img.verify()
        return Image.open(io.BytesIO(image_bytes)).convert("RGB")

    else:
        raise RuntimeError(f"Unsupported IMAGE_MODEL: {effective_model}")



def get_admin_db(appname: Optional[str]):
    """
    Returns (categories_collection, media_clicks_collection)
    based on appname
    """
    if appname == "collage-maker":
        collage_uri = os.getenv("COLLAGE_MAKER_DB_URL")
        if not collage_uri:
            raise RuntimeError("COLLAGE_MAKER_DB_URL not set")

        client = MongoClient(collage_uri)
        db = client["adminPanel"]
        return db.categories, db.media_clicks

    # DEFAULT (existing behavior)
    admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
    db = admin_client["adminPanel"]
    return db.categories, db.media_clicks

# ---------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------
@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "success": True,
        "message": "Polaroid,Kiddo,Makeup,Hairstyle API",
        "data": {
            "version": "1.0.0",
            "Product Name":"Beauty Camera - GlowCam AI Studio",
            "Released By" : "LogicGo Infotech"
        }
    }
    
@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),
    appname: 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}")

    # -------------------------
    # 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, media_clicks_col = get_admin_db(appname)
            # 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:
        # --------------------------------------------------
        # MODEL OVERRIDE BASED ON CATEGORY
        # --------------------------------------------------
        force_model = None
        
        if category_id == GEMINI_FORCE_CATEGORY_ID:
            force_model = "GEMINI"
        
        pil_output = run_image_generation(
            image1=pil_img1,
            image2=pil_img2,
            prompt=prompt,
            force_model="GEMINI" if category_id == GEMINI_FORCE_CATEGORY_ID else None
        )


    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,
            "appname": appname
        })
        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,
        "appname": appname
    })

    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}"
        )
    })

@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)