main.py

# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: Apache-2.0

import os
import logging
from fastapi import FastAPI, Request, Depends, HTTPException, status, Query
from fastapi.responses import HTMLResponse, RedirectResponse, PlainTextResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.sessions import SessionMiddleware
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.middleware.cors import CORSMiddleware
from api.endpoints import router as api_router
from api.auth import fastapi_users, auth_backend, google_oauth_client, github_oauth_client
from api.database import get_db, engine, Base
from api.models import User, UserRead, UserCreate, Conversation
from motor.motor_asyncio import AsyncIOMotorClient
from pydantic import BaseModel
from typing import List
from contextlib import asynccontextmanager
import uvicorn
import markdown2
from sqlalchemy.orm import Session
from pathlib import Path
from hashlib import md5
from datetime import datetime
import re
from httpx_oauth.exceptions import GetIdEmailError

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.info("Files in current dir: %s", os.listdir(os.getcwd()))

# Check environment variables
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    logger.error("HF_TOKEN is not set in environment variables.")
    raise ValueError("HF_TOKEN is required for Inference API.")

MONGO_URI = os.getenv("MONGODB_URI")
if not MONGO_URI:
    logger.error("MONGODB_URI is not set in environment variables.")
    raise ValueError("MONGODB_URI is required for MongoDB.")

JWT_SECRET = os.getenv("JWT_SECRET")
if not JWT_SECRET or len(JWT_SECRET) < 32:
    logger.error("JWT_SECRET is not set or too short.")
    raise ValueError("JWT_SECRET is required (at least 32 characters).")

# MongoDB setup
client = AsyncIOMotorClient(MONGO_URI)
mongo_db = client["hager"]
session_message_counts = mongo_db["session_message_counts"]

# Create MongoDB index
async def setup_mongo_index():
    await session_message_counts.create_index("session_id", unique=True)

# Jinja2 setup
templates = Jinja2Templates(directory="templates")
templates.env.filters['markdown'] = lambda text: markdown2.markdown(text)

# Pydantic model for blog posts
class BlogPost(BaseModel):
    id: str
    title: str
    content: str
    author: str
    date: str
    created_at: str

# Application settings
QUEUE_SIZE = int(os.getenv("QUEUE_SIZE", 80))
CONCURRENCY_LIMIT = int(os.getenv("CONCURRENCY_LIMIT", 20))

# Initialize FastAPI app
@asynccontextmanager
async def lifespan(app: FastAPI):
    await setup_mongo_index()
    yield

app = FastAPI(title="MGZon Chatbot API", lifespan=lifespan)

# Add SessionMiddleware
app.add_middleware(SessionMiddleware, secret_key=JWT_SECRET)

# Create SQLAlchemy tables
Base.metadata.create_all(bind=engine)

# Mount static files
os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

# CORS setup
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://mgzon-mgzon-app.hf.space",
        "http://localhost:7860",
    ],
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Accept", "Content-Type", "Authorization"],
)

# Include routers
app.include_router(
    fastapi_users.get_auth_router(auth_backend),
    prefix="/auth/jwt",
    tags=["auth"],
)
app.include_router(
    fastapi_users.get_register_router(UserRead, UserCreate),
    prefix="/auth",
    tags=["auth"],
)
app.include_router(
    fastapi_users.get_users_router(UserRead, UserCreate),
    prefix="/users",
    tags=["users"],
)
app.include_router(
    fastapi_users.get_oauth_router(
        google_oauth_client,
        auth_backend,
        JWT_SECRET,
        redirect_url="https://mgzon-mgzon-app.hf.space/auth/google/callback"
    ),
    prefix="/auth/google",
    tags=["auth"],
)
app.include_router(
    fastapi_users.get_oauth_router(
        github_oauth_client,
        auth_backend,
        JWT_SECRET,
        redirect_url="https://mgzon-mgzon-app.hf.space/auth/github/callback"
    ),
    prefix="/auth/github",
    tags=["auth"],
)
app.include_router(api_router)

# Debug routes endpoint
@app.get("/debug/routes", response_class=PlainTextResponse)
async def debug_routes():
    routes = []
    for route in app.routes:
        methods = getattr(route, "methods", [])
        path = getattr(route, "path", "Unknown")
        routes.append(f"{methods} {path}")
    return "\n".join(sorted(routes))

# OAuth callbacks
@app.get("/auth/google/callback", response_class=RedirectResponse)
async def google_oauth_callback(
    request: Request,
    token: str = Query(...),
    state: str = Query(...),
    db: Session = Depends(get_db)
):
    try:
        logger.info("Processing Google OAuth callback")
        # Exchange code for access token
        token_data = await google_oauth_client.get_access_token(token, "https://mgzon-mgzon-app.hf.space/auth/google/callback")
        logger.info(f"Google OAuth token received: {token_data}")
        # Get user info
        user_info = await google_oauth_client.get_id_email(token_data["access_token"])
        logger.info(f"Google user info: {user_info}")
        # Create or update user
        user_manager = fastapi_users.user_manager
        user = await user_manager.oauth_callback(
            oauth_name="google",
            access_token=token_data["access_token"],
            account_id=user_info["id"],
            account_email=user_info["email"],
            expires_at=token_data.get("expires_at"),
            refresh_token=token_data.get("refresh_token"),
            request=request,
            db=db
        )
        logger.info("Google OAuth user processed, creating session")
        # Create JWT token
        token = await auth_backend.get_login_response(user, request)
        logger.info("Google OAuth callback processed, redirecting to /chat")
        response = RedirectResponse(url="/chat", status_code=302)
        response.set_cookie("Authorization", f"Bearer {token.access_token}", httponly=True)
        return response
    except Exception as e:
        logger.error(f"Google OAuth callback error: {str(e)}")
        return RedirectResponse(url=f"/login?error=Google%20OAuth%20failed:%20{str(e)}", status_code=302)

@app.get("/auth/github/callback", response_class=RedirectResponse)
async def github_oauth_callback(
    request: Request,
    token: str = Query(...),
    state: str = Query(...),
    db: Session = Depends(get_db)
):
    try:
        logger.info("Processing GitHub OAuth callback")
        # Exchange code for access token
        token_data = await github_oauth_client.get_access_token(token, "https://mgzon-mgzon-app.hf.space/auth/github/callback")
        logger.info(f"GitHub OAuth token received: {token_data}")
        # Get user info
        user_info = await github_oauth_client.get_id_email(token_data["access_token"])
        logger.info(f"GitHub user info: {user_info}")
        # Create or update user
        user_manager = fastapi_users.user_manager
        user = await user_manager.oauth_callback(
            oauth_name="github",
            access_token=token_data["access_token"],
            account_id=user_info["id"],
            account_email=user_info["email"],
            expires_at=token_data.get("expires_at"),
            refresh_token=token_data.get("refresh_token"),
            request=request,
            db=db
        )
        logger.info("GitHub OAuth user processed, creating session")
        # Create JWT token
        token = await auth_backend.get_login_response(user, request)
        logger.info("GitHub OAuth callback processed, redirecting to /chat")
        response = RedirectResponse(url="/chat", status_code=302)
        response.set_cookie("Authorization", f"Bearer {token.access_token}", httponly=True)
        return response
    except Exception as e:
        logger.error(f"GitHub OAuth callback error: {str(e)}")
        return RedirectResponse(url=f"/login?error=GitHub%20OAuth%20failed:%20{str(e)}", status_code=302)

# Custom middleware for 404 and 500 errors
class NotFoundMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        try:
            response = await call_next(request)
            if response.status_code == 404:
                logger.warning(f"404 Not Found: {request.url}")
                return templates.TemplateResponse("404.html", {"request": request}, status_code=404)
            return response
        except Exception as e:
            logger.exception(f"Error processing request {request.url}: {e}")
            return templates.TemplateResponse("500.html", {"request": request, "error": str(e)}, status_code=500)

app.add_middleware(NotFoundMiddleware)

# OAuth error handler
@app.exception_handler(GetIdEmailError)
async def handle_oauth_error(request: Request, exc: GetIdEmailError):
    logger.error(f"OAuth error: {exc}")
    error_message = "Failed to authenticate with OAuth. Please try again or contact support."
    return RedirectResponse(
        url=f"/login?error={error_message}",
        status_code=302
    )

# Root endpoint
@app.get("/", response_class=HTMLResponse)
async def root(request: Request, user: User = Depends(current_active_user)):
    return templates.TemplateResponse("index.html", {"request": request, "user": user})

# Google verification
@app.get("/google97468ef1f6b6e804.html", response_class=PlainTextResponse)
async def google_verification():
    return "google-site-verification: google97468ef1f6b6e804.html"

# Login page
@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request, user: User = Depends(current_active_user)):
    if user:
        return RedirectResponse(url="/chat", status_code=302)
    return templates.TemplateResponse("login.html", {"request": request})

# Register page
@app.get("/register", response_class=HTMLResponse)
async def register_page(request: Request, user: User = Depends(current_active_user)):
    if user:
        return RedirectResponse(url="/chat", status_code=302)
    return templates.TemplateResponse("register.html", {"request": request})

# Chat page
@app.get("/chat", response_class=HTMLResponse)
async def chat(request: Request, user: User = Depends(current_active_user)):
    return templates.TemplateResponse("chat.html", {"request": request, "user": user})

# Specific conversation page
@app.get("/chat/{conversation_id}", response_class=HTMLResponse)
async def chat_conversation(
    request: Request,
    conversation_id: str,
    user: User = Depends(current_active_user),
    db: Session = Depends(get_db)
):
    if not user:
        return RedirectResponse(url="/login", status_code=302)
    conversation = db.query(Conversation).filter(
        Conversation.conversation_id == conversation_id,
        Conversation.user_id == user.id
    ).first()
    if not conversation:
        raise HTTPException(status_code=404, detail="Conversation not found")
    return templates.TemplateResponse(
        "chat.html",
        {
            "request": request,
            "user": user,
            "conversation_id": conversation.conversation_id,
            "conversation_title": conversation.title or "Untitled Conversation"
        }
    )

# About page
@app.get("/about", response_class=HTMLResponse)
async def about(request: Request, user: User = Depends(current_active_user)):
    return templates.TemplateResponse("about.html", {"request": request, "user": user})

# Serve static files
@app.get("/static/{path:path}")
async def serve_static(path: str):
    clean_path = re.sub(r'\?.*', '', path)
    file_path = Path("static") / clean_path
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found")
    cache_duration = 31536000 if not clean_path.endswith(('.js', '.css')) else 3600
    with open(file_path, "rb") as f:
        file_hash = md5(f.read()).hexdigest()
    headers = {
        "Cache-Control": f"public, max-age={cache_duration}",
        "ETag": file_hash,
        "Last-Modified": datetime.utcfromtimestamp(file_path.stat().st_mtime).strftime('%a, %d %b %Y %H:%M:%S GMT')
    }
    return FileResponse(file_path, headers=headers)

# Blog page
@app.get("/blog", response_class=HTMLResponse)
async def blog(request: Request, skip: int = Query(0, ge=0), limit: int = Query(10, ge=1, le=100)):
    posts = await mongo_db.blog_posts.find().skip(skip).limit(limit).to_list(limit)
    return templates.TemplateResponse("blog.html", {"request": request, "posts": posts})

# Individual blog post
@app.get("/blog/{post_id}", response_class=HTMLResponse)
async def blog_post(request: Request, post_id: str):
    post = await mongo_db.blog_posts.find_one({"id": post_id})
    if not post:
        raise HTTPException(status_code=404, detail="Post not found")
    return templates.TemplateResponse("blog_post.html", {"request": request, "post": post})

# Docs page
@app.get("/docs", response_class=HTMLResponse)
async def docs(request: Request):
    return templates.TemplateResponse("docs.html", {"request": request})

# Swagger UI
@app.get("/swagger", response_class=HTMLResponse)
async def swagger_ui():
    return get_swagger_ui_html(openapi_url="/openapi.json", title="MGZon API Documentation")

# Sitemap
@app.get("/sitemap.xml", response_class=PlainTextResponse)
async def sitemap():
    posts = await mongo_db.blog_posts.find().to_list(100)
    current_date = datetime.utcnow().strftime('%Y-%m-%d')
    xml = '<?xml version="1.0" encoding="UTF-8"?>\n'
    xml += '<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>daily</changefreq>\n'
    xml += '    <priority>1.0</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/chat</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>daily</changefreq>\n'
    xml += '    <priority>0.8</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/about</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>weekly</changefreq>\n'
    xml += '    <priority>0.7</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/login</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>weekly</changefreq>\n'
    xml += '    <priority>0.8</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/register</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>weekly</changefreq>\n'
    xml += '    <priority>0.8</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/docs</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>weekly</changefreq>\n'
    xml += '    <priority>0.9</priority>\n'
    xml += '  </url>\n'
    xml += '  <url>\n'
    xml += '    <loc>https://mgzon-mgzon-app.hf.space/blog</loc>\n'
    xml += f'    <lastmod>{current_date}</lastmod>\n'
    xml += '    <changefreq>daily</changefreq>\n'
    xml += '    <priority>0.9</priority>\n'
    xml += '  </url>\n'
    for post in posts:
        xml += '  <url>\n'
        xml += f'    <loc>https://mgzon-mgzon-app.hf.space/blog/{post["id"]}</loc>\n'
        xml += f'    <lastmod>{post["date"]}</lastmod>\n'
        xml += '    <changefreq>weekly</changefreq>\n'
        xml += '    <priority>0.9</priority>\n'
        xml += '  </url>\n'
    xml += '</urlset>'
    return xml

# Redirect /gradio to /chat
@app.get("/gradio", response_class=RedirectResponse)
async def launch_chatbot():
    return RedirectResponse(url="/chat", status_code=302)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))