main.py

from fastapi import FastAPI, File, UploadFile, Form, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
import time
import config
from models import APIResponse, VideoGenerationResponse
from services.pixverse_service import pixverse_service
from services.storage_service import storage_service
from services.firebase_service import firebase_service
from services.mongodb_service import save_media_click, save_request_log, get_category_by_prompt, get_selfi_short_client

app = FastAPI(
    title="PixVerse Video Generation API",
    description="API for generating videos from images using PixVerse AI",
    version="1.0.0"
)

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


def verify_hf_token(authorization: str) -> bool:
    """Verify Hugging Face token from Authorization header (optional if HF_TOKEN not set)"""
    expected_token = config.HF_TOKEN
    
    # If HF_TOKEN is not configured, skip verification (optional)
    if not expected_token:
        print(f"[Auth] HF_TOKEN not configured - skipping authorization check")
        return True
    
    if not authorization:
        print(f"[Auth] No Authorization header provided")
        return False
    
    token = authorization.replace("Bearer ", "").strip()
    
    # Debug logging
    print(f"[Auth] Received token: {token[:10]}... (length: {len(token)})")
    print(f"[Auth] Expected token: {expected_token[:10]}... (length: {len(expected_token)})")
    print(f"[Auth] Tokens match: {token == expected_token}")
    
    return token == expected_token


@app.get("/health")
async def health_check():
    """Health check endpoint - checks server and MongoDB status"""
    start_time = time.time()
    
    # Check MongoDB connection
    mongodb_status = "disconnected"
    try:
        mongo_client = get_selfi_short_client()
        if mongo_client:
            mongo_client.admin.command('ping')
            mongodb_status = "connected"
    except Exception as e:
        print(f"[Health] MongoDB check failed: {e}")
        mongodb_status = "disconnected"
    
    # Calculate response time
    response_time = round(time.time() - start_time, 3)
    
    return APIResponse(
        success=True,
        message="Health check completed",
        data={
            "server": "running",
            "mongodb": mongodb_status,
            "gemini": "available",
            "qwen": "available",
            "response_time": f"{response_time}s"
        }
    )


@app.post("/get-firebase-user")
async def get_firebase_user(
    firebase_token: str = Form(..., description="Firebase ID token"),
    authorization: Optional[str] = Header(None, description="HF Token as Bearer token")
):
    """
    Get Firebase user ID from authentication token
    
    - Verifies the Firebase ID token
    - Returns user ID (uid) and email from the token
    """
    try:
        # Step 1: Verify HF Token
        if not verify_hf_token(authorization):
            raise HTTPException(
                status_code=401,
                detail="Invalid or missing Hugging Face token"
            )
        
        # Step 2: Verify Firebase token and get user info
        try:
            user_info = firebase_service.verify_token(firebase_token)
            return APIResponse(
                success=True,
                message="Firebase user verified successfully",
                data={
                    "user_id": user_info.get("uid"),
                    "email": user_info.get("email"),
                    "verified": user_info.get("verified", True)
                }
            )
        except Exception as e:
            raise HTTPException(
                status_code=401,
                detail=f"Firebase authentication failed: {str(e)}"
            )
    
    except HTTPException:
        raise
    except Exception as e:
        return APIResponse(
            success=False,
            message=str(e),
            data=None
        )


@app.post("/generate-video")
async def generate_video(
    image: UploadFile = File(..., description="User image file (JPG, PNG, WebP)"),
    user_id: Optional[str] = Form(None, description="User ID (ObjectId string or integer)"),
    prompt_text: str = Form(..., description="Prompt text for video generation (must contain keywords: kiss, hug, bf/boyfriend, gf/girlfriend)"),
    firebase_id_token: str = Form(..., description="Firebase ID token"),
    appname: Optional[str] = Form(None, description="App name for database routing (collage-maker, AI-Enhancer, etc.)"),
    authorization: Optional[str] = Header(None, description="HF Token as Bearer token")
):
    """
    Generate video from image using PixVerse AI
    
    - Uploads user image to Digital Ocean Spaces (milestone/valentine/source/)
    - Generates video using PixVerse API
    - Stores result video in Digital Ocean Spaces (milestone/valentine/results/)
    - Automatically matches prompt_text to one of 4 categories
    - Tracks media clicks in MongoDB based on appname
    - Logs request to MongoDB Logs database
    - Returns URLs for both source image and generated video
    
    Category Matching (based on prompt_text keywords):
    - "kiss" → Ai Kiss Video
    - "hug" → Ai Hug Video
    - "bf" or "boyfriend" → Ai BF
    - "gf" or "girlfriend" → Ai GF
    
    Database Routing (based on appname):
    - "collage-maker" → collage-maker MongoDB → adminPanel DB → media_clicks collection
    - "AI-Enhancer" → AI-Enhancer MongoDB → test DB → media_clicks collection
    - (default/other) → Selfie Short MongoDB → adminPanel DB → media_clicks collection
    """
    start_time = time.time()
    
    try:
        # Step 1: Verify HF Token
        if not verify_hf_token(authorization):
            raise HTTPException(
                status_code=401,
                detail="Invalid or missing Hugging Face token"
            )
        
        # Step 2: Verify Firebase authentication
        try:
            firebase_service.verify_token(firebase_id_token)
        except Exception as e:
            raise HTTPException(
                status_code=401,
                detail=f"Firebase authentication failed: {str(e)}"
            )
        
        # Step 3: Match prompt_text to category (Pixverse Kiss Video / Pixverse Hug Video)
        category_id, category_name = get_category_by_prompt(prompt_text, appname)
        print(f"[API] Matched prompt to category - category_id: {category_id}, category_name: {category_name}")
        
        # Step 4: Validate image file
        if not image.content_type in ["image/jpeg", "image/png", "image/webp"]:
            raise HTTPException(
                status_code=400,
                detail="Invalid image format. Supported formats: JPG, PNG, WebP"
            )
        
        # Step 5: Read image content
        image_content = await image.read()
        
        # Check file size (max 20MB)
        if len(image_content) > 20 * 1024 * 1024:
            raise HTTPException(
                status_code=400,
                detail="Image file size exceeds 20MB limit"
            )
        
        # Step 6: Upload source image to Digital Ocean Spaces
        source_image_url = storage_service.upload_source_image(
            file_content=image_content,
            user_id=user_id or "anonymous",
            category_id=category_id or "default",
            content_type=image.content_type
        )
        
        # Step 7: Upload image to PixVerse
        upload_response = await pixverse_service.upload_image(
            image_content=image_content,
            content_type=image.content_type
        )
        
        if upload_response.get("ErrCode") != 0:
            raise HTTPException(
                status_code=500,
                detail=f"PixVerse image upload failed: {upload_response.get('ErrMsg')}"
            )
        
        img_id = upload_response.get("Resp", {}).get("img_id")
        if not img_id:
            raise HTTPException(
                status_code=500,
                detail="Failed to get image ID from PixVerse"
            )
        
        # Step 8: Generate video using PixVerse
        generate_response = await pixverse_service.generate_video(
            img_id=img_id,
            prompt=prompt_text
        )
        
        if generate_response.get("ErrCode") != 0:
            raise HTTPException(
                status_code=500,
                detail=f"PixVerse video generation failed: {generate_response.get('ErrMsg')}"
            )
        
        video_id = generate_response.get("Resp", {}).get("video_id")
        if not video_id:
            raise HTTPException(
                status_code=500,
                detail="Failed to get video ID from PixVerse"
            )
        
        # Step 9: Wait for video generation to complete
        video_result = await pixverse_service.wait_for_video(video_id)
        
        pixverse_video_url = video_result.get("Resp", {}).get("url")
        if not pixverse_video_url:
            raise HTTPException(
                status_code=500,
                detail="Failed to get video URL from PixVerse"
            )
        
        # Step 10: Download video from PixVerse
        video_content = await pixverse_service.download_video(pixverse_video_url)
        
        # Step 11: Upload result video to Digital Ocean Spaces
        result_video_url = storage_service.upload_result_video(
            video_content=video_content,
            user_id=user_id or "anonymous",
            category_id=category_id or "default"
        )
        
        # Step 12: Save media click to MongoDB (only if user_id and category_id are matched)
        if category_id and user_id is not None:
            print(f"[API] Saving media click - category_id: {category_id}, user_id: {user_id}, appname: {appname}")
            save_media_click(user_id=user_id, category_id=category_id, appname=appname)
        elif category_id and user_id is None:
            print(f"[API] Skipping media_clicks logging - user_id not provided")
        else:
            print(f"[API] Skipping media_clicks logging - no matching category found in prompt")
        
        # Step 13: Log request to MongoDB Logs
        response_time = time.time() - start_time
        # Use category_name from prompt matching (already retrieved in Step 3)
        subcategory_name = category_name or "Valentine Video Generation"
        save_request_log(
            user_id=user_id,
            subcategory=subcategory_name or "Valentine Video Generation",
            endpoint="/generate-video",
            status="success",
            response_time=response_time,
            model="PixVerse",
            appname=appname,
            error=None
        )
        
        # Return success response
        return APIResponse(
            success=True,
            message="Video created successfully",
            data=VideoGenerationResponse(
                result_url=result_video_url,
                source_image=source_image_url,
                video_id=video_id
            ).model_dump()
        )
        
    except HTTPException as he:
        # Log error request
        response_time = time.time() - start_time
        # Use category_name from prompt matching if available
        log_subcategory = category_name if 'category_name' in dir() and category_name else "Valentine Video Generation"
        save_request_log(
            user_id=user_id if 'user_id' in dir() else None,
            subcategory=log_subcategory,
            endpoint="/generate-video",
            status="failure",
            response_time=response_time,
            model="PixVerse",
            appname=appname if 'appname' in dir() else None,
            error=str(he.detail)
        )
        raise
    except Exception as e:
        # Log error request
        response_time = time.time() - start_time
        # Use category_name from prompt matching if available
        log_subcategory = category_name if 'category_name' in dir() and category_name else "Valentine Video Generation"
        save_request_log(
            user_id=user_id if 'user_id' in dir() else None,
            subcategory=log_subcategory,
            endpoint="/generate-video",
            status="failure",
            response_time=response_time,
            model="PixVerse",
            appname=appname if 'appname' in dir() else None,
            error=str(e)
        )
        return APIResponse(
            success=False,
            message=f"Error processing video: {e}",
            data=None
        )


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