api/api_server.py

"""
REST API server for BackgroundFX Pro.
Provides HTTP endpoints for all processing functionality.
"""

from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks, Depends, status
from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field, validator
from typing import Dict, List, Optional, Union, Any
from enum import Enum
import asyncio
import aiofiles
from pathlib import Path
import tempfile
import shutil
import uuid
import time
from datetime import datetime, timedelta
import jwt
import cv2
import numpy as np
import io
import base64
from concurrent.futures import ThreadPoolExecutor
import redis
from contextlib import asynccontextmanager

from ..utils.logger import setup_logger
from .pipeline import ProcessingPipeline, PipelineConfig, ProcessingMode
from .video_processor import VideoProcessorAPI, StreamConfig, VideoStreamMode
from .batch_processor import BatchProcessor, BatchConfig, BatchItem, BatchPriority

logger = setup_logger(__name__)


# ============================================================================
# Configuration and Models
# ============================================================================

class ServerConfig:
    """Server configuration."""
    HOST: str = "0.0.0.0"
    PORT: int = 8000
    UPLOAD_DIR: str = "uploads"
    OUTPUT_DIR: str = "outputs"
    TEMP_DIR: str = "temp"
    MAX_UPLOAD_SIZE: int = 500 * 1024 * 1024  # 500MB
    ALLOWED_EXTENSIONS: List[str] = [".jpg", ".jpeg", ".png", ".mp4", ".avi", ".mov"]
    
    # Security
    SECRET_KEY: str = "your-secret-key-change-in-production"
    ALGORITHM: str = "HS256"
    ACCESS_TOKEN_EXPIRE_MINUTES: int = 30
    
    # Redis cache
    REDIS_URL: str = "redis://localhost:6379"
    CACHE_TTL: int = 3600  # 1 hour
    
    # Rate limiting
    RATE_LIMIT_REQUESTS: int = 100
    RATE_LIMIT_WINDOW: int = 60  # seconds
    
    # Processing
    MAX_WORKERS: int = 4
    ENABLE_GPU: bool = True


config = ServerConfig()


# ============================================================================
# Pydantic Models
# ============================================================================

class BackgroundType(str, Enum):
    """Background types."""
    BLUR = "blur"
    OFFICE = "office"
    GRADIENT = "gradient"
    NATURE = "nature"
    CUSTOM = "custom"
    NONE = "none"


class QualityPreset(str, Enum):
    """Quality presets."""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    ULTRA = "ultra"


class ProcessingRequest(BaseModel):
    """Base processing request."""
    background: BackgroundType = BackgroundType.BLUR
    background_url: Optional[str] = None
    quality: QualityPreset = QualityPreset.HIGH
    preserve_original: bool = False
    
    class Config:
        schema_extra = {
            "example": {
                "background": "office",
                "quality": "high",
                "preserve_original": False
            }
        }


class ImageProcessingRequest(ProcessingRequest):
    """Image processing request."""
    resize: Optional[tuple[int, int]] = None
    apply_effects: List[str] = Field(default_factory=list)
    output_format: str = "png"


class VideoProcessingRequest(ProcessingRequest):
    """Video processing request."""
    start_time: Optional[float] = None
    end_time: Optional[float] = None
    fps: Optional[float] = None
    resolution: Optional[tuple[int, int]] = None
    codec: str = "h264"


class BatchProcessingRequest(BaseModel):
    """Batch processing request."""
    items: List[Dict[str, Any]]
    parallel: bool = True
    priority: str = "normal"
    callback_url: Optional[str] = None


class StreamingRequest(BaseModel):
    """Streaming request."""
    source: str
    stream_type: str = "webcam"
    output_format: str = "hls"
    quality: QualityPreset = QualityPreset.MEDIUM


class ProcessingResponse(BaseModel):
    """Processing response."""
    job_id: str
    status: str
    progress: float = 0.0
    message: Optional[str] = None
    result_url: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.now)
    completed_at: Optional[datetime] = None


class JobStatus(BaseModel):
    """Job status response."""
    job_id: str
    status: str
    progress: float
    current_stage: Optional[str] = None
    time_elapsed: float
    time_remaining: Optional[float] = None
    errors: List[str] = Field(default_factory=list)


# ============================================================================
# Job Management
# ============================================================================

class JobManager:
    """Manage processing jobs."""
    
    def __init__(self):
        self.jobs: Dict[str, ProcessingResponse] = {}
        self.executor = ThreadPoolExecutor(max_workers=config.MAX_WORKERS)
        self.redis_client = None
        try:
            self.redis_client = redis.from_url(config.REDIS_URL)
        except:
            logger.warning("Redis not available, using in-memory storage")
    
    def create_job(self) -> str:
        """Create new job ID."""
        job_id = str(uuid.uuid4())
        self.jobs[job_id] = ProcessingResponse(
            job_id=job_id,
            status="pending"
        )
        return job_id
    
    def update_job(self, job_id: str, **kwargs):
        """Update job status."""
        if job_id in self.jobs:
            for key, value in kwargs.items():
                if hasattr(self.jobs[job_id], key):
                    setattr(self.jobs[job_id], key, value)
            
            # Store in Redis if available
            if self.redis_client:
                try:
                    self.redis_client.setex(
                        f"job:{job_id}",
                        config.CACHE_TTL,
                        self.jobs[job_id].json()
                    )
                except:
                    pass
    
    def get_job(self, job_id: str) -> Optional[ProcessingResponse]:
        """Get job status."""
        # Check memory first
        if job_id in self.jobs:
            return self.jobs[job_id]
        
        # Check Redis
        if self.redis_client:
            try:
                data = self.redis_client.get(f"job:{job_id}")
                if data:
                    return ProcessingResponse.parse_raw(data)
            except:
                pass
        
        return None


# ============================================================================
# FastAPI Application
# ============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager."""
    # Startup
    logger.info("Starting BackgroundFX Pro API Server")
    
    # Create directories
    for dir_path in [config.UPLOAD_DIR, config.OUTPUT_DIR, config.TEMP_DIR]:
        Path(dir_path).mkdir(parents=True, exist_ok=True)
    
    # Initialize processors
    app.state.pipeline = ProcessingPipeline(
        PipelineConfig(use_gpu=config.ENABLE_GPU)
    )
    app.state.video_processor = VideoProcessorAPI()
    app.state.batch_processor = BatchProcessor()
    app.state.job_manager = JobManager()
    
    yield
    
    # Shutdown
    logger.info("Shutting down BackgroundFX Pro API Server")
    app.state.pipeline.shutdown()
    app.state.video_processor.cleanup()
    app.state.batch_processor.cleanup()


app = FastAPI(
    title="BackgroundFX Pro API",
    description="Professional background removal and replacement API",
    version="1.0.0",
    lifespan=lifespan
)

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


# ============================================================================
# Authentication
# ============================================================================

security = HTTPBearer()


def create_access_token(data: dict) -> str:
    """Create JWT access token."""
    to_encode = data.copy()
    expire = datetime.utcnow() + timedelta(minutes=config.ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)


def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
    """Verify JWT token."""
    token = credentials.credentials
    try:
        payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid authentication credentials",
            )
        return username
    except jwt.PyJWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication credentials",
        )


# ============================================================================
# Health and Status Endpoints
# ============================================================================

@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "name": "BackgroundFX Pro API",
        "version": "1.0.0",
        "status": "running",
        "endpoints": {
            "health": "/health",
            "docs": "/docs",
            "process_image": "/api/v1/process/image",
            "process_video": "/api/v1/process/video",
            "batch": "/api/v1/batch",
            "stream": "/api/v1/stream"
        }
    }


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "services": {
            "pipeline": "ready",
            "video_processor": "ready",
            "batch_processor": "ready",
            "redis": "connected" if app.state.job_manager.redis_client else "disconnected"
        }
    }


@app.get("/api/v1/stats")
async def get_statistics(current_user: str = Depends(verify_token)):
    """Get processing statistics."""
    return {
        "pipeline": app.state.pipeline.get_statistics(),
        "video": app.state.video_processor.get_stats(),
        "batch": app.state.batch_processor.get_status()
    }


# ============================================================================
# Image Processing Endpoints
# ============================================================================

@app.post("/api/v1/process/image", response_model=ProcessingResponse)
async def process_image(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    request: ImageProcessingRequest = Depends(),
    current_user: str = Depends(verify_token)
):
    """Process a single image."""
    
    # Validate file
    if not file.filename.lower().endswith(tuple(config.ALLOWED_EXTENSIONS)):
        raise HTTPException(400, "Invalid file format")
    
    if file.size > config.MAX_UPLOAD_SIZE:
        raise HTTPException(413, "File too large")
    
    # Create job
    job_id = app.state.job_manager.create_job()
    
    # Save uploaded file
    upload_path = Path(config.UPLOAD_DIR) / f"{job_id}_{file.filename}"
    async with aiofiles.open(upload_path, 'wb') as f:
        content = await file.read()
        await f.write(content)
    
    # Process in background
    background_tasks.add_task(
        process_image_task,
        app.state,
        job_id,
        str(upload_path),
        request
    )
    
    return ProcessingResponse(
        job_id=job_id,
        status="processing",
        message="Image processing started"
    )


async def process_image_task(app_state, job_id: str, input_path: str, request: ImageProcessingRequest):
    """Background task for image processing."""
    try:
        # Update job status
        app_state.job_manager.update_job(job_id, status="processing", progress=0.1)
        
        # Load image
        image = cv2.imread(input_path)
        
        # Prepare background
        background = None
        if request.background == BackgroundType.CUSTOM and request.background_url:
            # Download custom background
            # ... implementation ...
            pass
        elif request.background != BackgroundType.NONE:
            background = request.background.value
        
        # Configure pipeline
        config = PipelineConfig(
            quality_preset=request.quality.value,
            apply_effects=request.apply_effects
        )
        
        # Process image
        result = app_state.pipeline.process_image(image, background)
        
        if result.success:
            # Save output
            output_filename = f"{job_id}_output.{request.output_format}"
            output_path = Path(config.OUTPUT_DIR) / output_filename
            cv2.imwrite(str(output_path), result.output_image)
            
            # Update job
            app_state.job_manager.update_job(
                job_id,
                status="completed",
                progress=1.0,
                result_url=f"/api/v1/download/{output_filename}",
                completed_at=datetime.now(),
                metadata={
                    "quality_score": result.quality_score,
                    "processing_time": result.processing_time
                }
            )
        else:
            app_state.job_manager.update_job(
                job_id,
                status="failed",
                message="Processing failed"
            )
            
    except Exception as e:
        logger.error(f"Image processing failed for job {job_id}: {e}")
        app_state.job_manager.update_job(
            job_id,
            status="failed",
            message=str(e)
        )


# ============================================================================
# Video Processing Endpoints
# ============================================================================

@app.post("/api/v1/process/video", response_model=ProcessingResponse)
async def process_video(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    request: VideoProcessingRequest = Depends(),
    current_user: str = Depends(verify_token)
):
    """Process a video file."""
    
    # Validate file
    if not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
        raise HTTPException(400, "Invalid video format")
    
    # Create job
    job_id = app_state.job_manager.create_job()
    
    # Save uploaded file
    upload_path = Path(config.UPLOAD_DIR) / f"{job_id}_{file.filename}"
    async with aiofiles.open(upload_path, 'wb') as f:
        content = await file.read()
        await f.write(content)
    
    # Process in background
    background_tasks.add_task(
        process_video_task,
        app.state,
        job_id,
        str(upload_path),
        request
    )
    
    return ProcessingResponse(
        job_id=job_id,
        status="processing",
        message="Video processing started"
    )


async def process_video_task(app_state, job_id: str, input_path: str, request: VideoProcessingRequest):
    """Background task for video processing."""
    try:
        # Progress callback
        def progress_callback(progress: float, info: Dict):
            app_state.job_manager.update_job(
                job_id,
                progress=progress,
                metadata=info
            )
        
        # Process video
        output_path = Path(config.OUTPUT_DIR) / f"{job_id}_output.mp4"
        
        stats = await app_state.video_processor.process_video_async(
            input_path,
            str(output_path),
            background=request.background.value if request.background != BackgroundType.NONE else None,
            progress_callback=progress_callback
        )
        
        # Update job
        app_state.job_manager.update_job(
            job_id,
            status="completed",
            progress=1.0,
            result_url=f"/api/v1/download/{output_path.name}",
            completed_at=datetime.now(),
            metadata={
                "frames_processed": stats.frames_processed,
                "processing_fps": stats.processing_fps,
                "avg_quality": stats.avg_quality_score
            }
        )
        
    except Exception as e:
        logger.error(f"Video processing failed for job {job_id}: {e}")
        app_state.job_manager.update_job(
            job_id,
            status="failed",
            message=str(e)
        )


# ============================================================================
# Batch Processing Endpoints
# ============================================================================

@app.post("/api/v1/batch", response_model=ProcessingResponse)
async def process_batch(
    background_tasks: BackgroundTasks,
    request: BatchProcessingRequest,
    current_user: str = Depends(verify_token)
):
    """Process multiple files in batch."""
    
    # Create job
    job_id = app.state.job_manager.create_job()
    
    # Process in background
    background_tasks.add_task(
        process_batch_task,
        app.state,
        job_id,
        request
    )
    
    return ProcessingResponse(
        job_id=job_id,
        status="processing",
        message=f"Batch processing started for {len(request.items)} items"
    )


async def process_batch_task(app_state, job_id: str, request: BatchProcessingRequest):
    """Background task for batch processing."""
    try:
        # Convert request items to BatchItems
        batch_items = []
        for item_data in request.items:
            batch_item = BatchItem(
                id=item_data.get('id', str(uuid.uuid4())),
                input_path=item_data['input_path'],
                output_path=item_data['output_path'],
                file_type=item_data.get('file_type', 'image'),
                priority=BatchPriority[request.priority.upper()],
                background=item_data.get('background')
            )
            batch_items.append(batch_item)
        
        # Progress callback
        def progress_callback(progress: float, info: Dict):
            app_state.job_manager.update_job(
                job_id,
                progress=progress,
                metadata=info
            )
        
        # Configure batch processor
        batch_config = BatchConfig(
            progress_callback=progress_callback,
            max_workers=config.MAX_WORKERS if request.parallel else 1
        )
        
        processor = BatchProcessor(batch_config)
        report = processor.process_batch(batch_items)
        
        # Update job
        app_state.job_manager.update_job(
            job_id,
            status="completed",
            progress=1.0,
            completed_at=datetime.now(),
            metadata={
                "total_items": report.total_items,
                "successful_items": report.successful_items,
                "failed_items": report.failed_items,
                "avg_quality": report.quality_metrics.get('avg_quality', 0)
            }
        )
        
        # Callback if provided
        if request.callback_url:
            # Send completion callback
            # ... implementation ...
            pass
            
    except Exception as e:
        logger.error(f"Batch processing failed for job {job_id}: {e}")
        app_state.job_manager.update_job(
            job_id,
            status="failed",
            message=str(e)
        )


# ============================================================================
# Streaming Endpoints
# ============================================================================

@app.post("/api/v1/stream/start")
async def start_stream(
    request: StreamingRequest,
    current_user: str = Depends(verify_token)
):
    """Start a streaming session."""
    
    # Configure streaming
    stream_config = StreamConfig(
        source=request.source,
        stream_mode=VideoStreamMode[request.stream_type.upper()],
        output_format=request.output_format,
        output_path=f"{config.OUTPUT_DIR}/stream_{uuid.uuid4()}"
    )
    
    # Start streaming
    success = app.state.video_processor.start_stream_processing(
        stream_config,
        background=None  # Configure as needed
    )
    
    if success:
        return {
            "status": "streaming",
            "stream_url": f"/api/v1/stream/live/{stream_config.output_path}",
            "message": "Streaming started"
        }
    else:
        raise HTTPException(500, "Failed to start streaming")


@app.get("/api/v1/stream/stop")
async def stop_stream(current_user: str = Depends(verify_token)):
    """Stop streaming session."""
    app.state.video_processor.stop_stream_processing()
    return {"status": "stopped", "message": "Streaming stopped"}


@app.get("/api/v1/stream/preview")
async def get_stream_preview(current_user: str = Depends(verify_token)):
    """Get stream preview frame."""
    frame = app.state.video_processor.get_preview_frame()
    
    if frame is not None:
        # Convert to JPEG
        _, buffer = cv2.imencode('.jpg', frame)
        return StreamingResponse(
            io.BytesIO(buffer),
            media_type="image/jpeg"
        )
    else:
        raise HTTPException(404, "No preview available")


# ============================================================================
# Job Management Endpoints
# ============================================================================

@app.get("/api/v1/job/{job_id}", response_model=ProcessingResponse)
async def get_job_status(
    job_id: str,
    current_user: str = Depends(verify_token)
):
    """Get job status."""
    job = app.state.job_manager.get_job(job_id)
    
    if job:
        return job
    else:
        raise HTTPException(404, "Job not found")


@app.get("/api/v1/jobs")
async def list_jobs(
    current_user: str = Depends(verify_token),
    limit: int = 10,
    offset: int = 0
):
    """List recent jobs."""
    jobs = list(app.state.job_manager.jobs.values())
    return {
        "total": len(jobs),
        "jobs": jobs[offset:offset + limit]
    }


@app.delete("/api/v1/job/{job_id}")
async def cancel_job(
    job_id: str,
    current_user: str = Depends(verify_token)
):
    """Cancel a job."""
    # Implementation would depend on your cancellation mechanism
    app.state.job_manager.update_job(job_id, status="cancelled")
    return {"message": "Job cancelled"}


# ============================================================================
# Download Endpoints
# ============================================================================

@app.get("/api/v1/download/{filename}")
async def download_file(
    filename: str,
    current_user: str = Depends(verify_token)
):
    """Download processed file."""
    file_path = Path(config.OUTPUT_DIR) / filename
    
    if file_path.exists():
        return FileResponse(
            path=file_path,
            filename=filename,
            media_type='application/octet-stream'
        )
    else:
        raise HTTPException(404, "File not found")


# ============================================================================
# WebSocket for Real-time Updates
# ============================================================================

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws/job/{job_id}")
async def websocket_job_updates(websocket: WebSocket, job_id: str):
    """WebSocket for real-time job updates."""
    await websocket.accept()
    
    try:
        while True:
            # Get job status
            job = app.state.job_manager.get_job(job_id)
            
            if job:
                await websocket.send_json(job.dict())
                
                if job.status in ["completed", "failed", "cancelled"]:
                    break
            
            await asyncio.sleep(1)
            
    except WebSocketDisconnect:
        logger.info(f"WebSocket disconnected for job {job_id}")


# ============================================================================
# Run Server
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    
    uvicorn.run(
        app,
        host=config.HOST,
        port=config.PORT,
        log_level="info"
    )