OCR/backend/app/main.py

"""
Tool_OCR - FastAPI Application Entry Point (V2)
Main application setup with CORS, routes, and startup/shutdown events
"""

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import logging
from pathlib import Path

from app.core.config import settings

# Ensure log directory exists before configuring logging
Path(settings.log_file).parent.mkdir(parents=True, exist_ok=True)

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler(settings.log_file),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan events"""
    # Startup
    logger.info("Starting Tool_OCR V2 application...")

    # Ensure all directories exist
    settings.ensure_directories()
    logger.info("All directories created/verified")

    logger.info("Application startup complete")

    yield

    # Shutdown
    logger.info("Shutting down Tool_OCR application...")


# Create FastAPI application
app = FastAPI(
    title="Tool_OCR V2",
    description="OCR Processing System with External Authentication & Task Isolation",
    version="2.0.0",
    lifespan=lifespan,
)

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


# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint with GPU status"""
    from app.services.ocr_service import OCRService

    response = {
        "status": "healthy",
        "service": "Tool_OCR V2",
        "version": "2.0.0",
    }

    # Add GPU status information
    try:
        # Create temporary OCRService instance to get GPU status
        # In production, this should be a singleton service
        ocr_service = OCRService()
        gpu_status = ocr_service.get_gpu_status()

        response["gpu"] = {
            "available": gpu_status.get("gpu_available", False),
            "enabled": gpu_status.get("gpu_enabled", False),
            "device_name": gpu_status.get("device_name", "N/A"),
            "device_count": gpu_status.get("device_count", 0),
            "compute_capability": gpu_status.get("compute_capability", "N/A"),
        }

        # Add memory info if available
        if gpu_status.get("memory_total_mb"):
            response["gpu"]["memory"] = {
                "total_mb": round(gpu_status.get("memory_total_mb", 0), 2),
                "allocated_mb": round(gpu_status.get("memory_allocated_mb", 0), 2),
                "utilization_percent": round(gpu_status.get("memory_utilization", 0), 2),
            }

        # Add reason if GPU is not available
        if not gpu_status.get("gpu_available") and gpu_status.get("reason"):
            response["gpu"]["reason"] = gpu_status.get("reason")

    except Exception as e:
        logger.warning(f"Failed to get GPU status: {e}")
        response["gpu"] = {
            "available": False,
            "error": str(e),
        }

    return response


# Root endpoint
@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "Tool_OCR API V2 - External Authentication",
        "version": "2.0.0",
        "docs_url": "/docs",
        "health_check": "/health",
    }


# Include V2 API routers
from app.routers import auth, tasks, admin
from fastapi import UploadFile, File, Depends, HTTPException, status
from sqlalchemy.orm import Session
import hashlib

from app.core.deps import get_db, get_current_user
from app.models.user import User
from app.models.task import TaskFile
from app.schemas.task import UploadResponse, TaskStatusEnum
from app.services.task_service import task_service

app.include_router(auth.router)
app.include_router(tasks.router)
app.include_router(admin.router)


# File upload endpoint
@app.post("/api/v2/upload", response_model=UploadResponse, tags=["Upload"], summary="Upload file for OCR")
async def upload_file(
    file: UploadFile = File(..., description="File to upload (PNG, JPG, PDF, etc.)"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Upload a file for OCR processing

    Creates a new task and uploads the file

    - **file**: File to upload
    """
    try:
        # Validate file extension
        file_ext = Path(file.filename).suffix.lower().lstrip('.')
        if file_ext not in settings.allowed_extensions_list:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"File type .{file_ext} not allowed. Allowed types: {', '.join(settings.allowed_extensions_list)}"
            )

        # Read file content
        file_content = await file.read()
        file_size = len(file_content)

        # Calculate file hash
        file_hash = hashlib.sha256(file_content).hexdigest()

        # Create task
        task = task_service.create_task(
            db=db,
            user_id=current_user.id,
            filename=file.filename,
            file_type=file.content_type
        )

        # Save file to disk
        upload_dir = Path(settings.upload_dir)
        upload_dir.mkdir(parents=True, exist_ok=True)

        # Create unique filename using task_id
        unique_filename = f"{task.task_id}_{file.filename}"
        file_path = upload_dir / unique_filename

        # Write file
        with open(file_path, "wb") as f:
            f.write(file_content)

        # Create TaskFile record
        task_file = TaskFile(
            task_id=task.id,
            original_name=file.filename,
            stored_path=str(file_path),
            file_size=file_size,
            mime_type=file.content_type,
            file_hash=file_hash
        )
        db.add(task_file)
        db.commit()

        logger.info(f"Uploaded file {file.filename} ({file_size} bytes) for task {task.task_id}, user {current_user.email}")

        return {
            "task_id": task.task_id,
            "filename": file.filename,
            "file_size": file_size,
            "file_type": file.content_type or "application/octet-stream",
            "status": TaskStatusEnum.PENDING
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"Failed to upload file for user {current_user.id}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to upload file: {str(e)}"
        )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=settings.backend_port,
        reload=True,
        log_level=settings.log_level.lower(),
    )