backend/main.py

"""
Antigravity Notebook - Main FastAPI Application
Entry point for the backend API server.
"""

from fastapi import FastAPI, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn

from backend.config import settings
from backend.database import init_db, engine
from backend.routers import notebooks, sources, chat

# Create FastAPI app
app = FastAPI(
    title="Antigravity Notebook",
    description="NotebookLM clone using CLaRa-7B latent compression for infinite context reasoning",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# CORS middleware (for Streamlit frontend)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify your frontend URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(notebooks.router)
app.include_router(sources.router)
app.include_router(chat.router)


@app.on_event("startup")
async def startup_event():
    """Initialize database and services on startup"""
    print("\n" + "="*60)
    print("🚀 Starting Antigravity Notebook")
    print("="*60)

    # Initialize database
    print("\n📊 Initializing database...")
    try:
        init_db()
        print("✅ Database ready!")
    except Exception as e:
        print(f"❌ Database initialization failed: {e}")

    # Pre-load CLaRa model (optional - can be lazy loaded)
    print("\n🤖 Initializing CLaRa model...")
    try:
        from backend.models.clara import get_clara_model
        clara = get_clara_model()
        print("✅ CLaRa model loaded!")
    except Exception as e:
        print(f"⚠️  CLaRa model failed to load: {e}")
        print("   Model will be loaded on first request")

    print("\n" + "="*60)
    print(f"✅ Antigravity Notebook is ready!")
    print(f"📍 API: http://{settings.API_HOST}:{settings.API_PORT}")
    print(f"📚 Docs: http://{settings.API_HOST}:{settings.API_PORT}/docs")
    print("="*60 + "\n")


@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    print("\n👋 Shutting down Antigravity Notebook...")
    engine.dispose()


@app.get("/", tags=["root"])
def root():
    """Root endpoint - health check"""
    return {
        "service": "Antigravity Notebook",
        "version": "1.0.0",
        "status": "running",
        "description": "NotebookLM clone with CLaRa-7B latent compression",
        "docs": "/docs"
    }


@app.get("/health", tags=["root"])
def health_check():
    """Health check endpoint"""
    try:
        # Check database connection
        from backend.database import SessionLocal
        db = SessionLocal()
        db.execute("SELECT 1")
        db.close()

        return {
            "status": "healthy",
            "database": "connected",
            "api": "running"
        }
    except Exception as e:
        return JSONResponse(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            content={
                "status": "unhealthy",
                "error": str(e)
            }
        )


@app.get("/stats", tags=["root"])
def get_stats():
    """Get storage statistics"""
    from backend.services.storage import get_storage_service

    storage = get_storage_service()
    stats = storage.get_storage_stats()

    return {
        "storage": stats,
        "config": {
            "max_context_tokens": settings.MAX_CONTEXT_TOKENS,
            "compression_ratio": settings.COMPRESSION_RATIO,
            "model": settings.MODEL_NAME
        }
    }


def main():
    """Run the application"""
    uvicorn.run(
        "backend.main:app",
        host=settings.API_HOST,
        port=settings.API_PORT,
        reload=True,
        log_level="info"
    )


if __name__ == "__main__":
    main()