app.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import asyncpg
import os
import hashlib
import logging
import secrets
import time
from typing import Dict
from datetime import timedelta
from dotenv import load_dotenv
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from contextlib import asynccontextmanager

# Load environment variables from .env file
load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)

DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:admin@localhost/chatdb")
CHALLENGE_EXPIRY = 300  # 5 minutes
SESSION_EXPIRY = 86400  # 24 hours
MIDNIGHT_CLEANUP = os.getenv("MIDNIGHT_CLEANUP", "false").lower() == "true"

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    try:
        app.state.db = await asyncpg.create_pool(DATABASE_URL)
        logger.info("Database connected successfully.")
        
        # Start midnight cleanup scheduler if enabled
        if MIDNIGHT_CLEANUP:
            scheduler.add_job(
                delete_all_messages,
                CronTrigger(hour=0, minute=0),  # Run at midnight
                id="midnight_message_cleanup",
                name="Delete all chat messages at midnight",
                replace_existing=True
            )
            scheduler.start()
            logger.info("Midnight message cleanup scheduler enabled - messages will be deleted daily at 00:00")
        else:
            logger.info("Midnight message cleanup is disabled")
            
    except Exception as e:
        logger.error(f"Failed to connect to database: {e}")
        raise
    
    yield  # Application runs here
    
    # Shutdown
    if scheduler.running:
        scheduler.shutdown()
    await app.state.db.close()
    logger.info("Database connection closed.")

app = FastAPI(lifespan=lifespan)

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

app.mount("/static", StaticFiles(directory="static"), name="static")

# Challenge and session management
challenges: Dict[str, tuple[str, float]] = {}  # {challenge_id: (challenge, expiry_time)}
sessions: Dict[str, tuple[str, str, float]] = {}  # {session_token: (chat_id, username, expiry_time)}

# Scheduler for midnight cleanup
scheduler = AsyncIOScheduler()

async def delete_all_messages():
    """Delete all messages from chat_message table at midnight"""
    try:
        async with app.state.db.acquire() as conn:
            deleted = await conn.execute("DELETE FROM chat_message")
            logger.info(f"Midnight cleanup: Deleted all messages from chat_message table")
    except Exception as e:
        logger.error(f"Error during midnight message cleanup: {e}")

def cleanup_expired_challenges():
    """Remove expired challenges"""
    current_time = time.time()
    expired = [cid for cid, (_, exp) in challenges.items() if exp < current_time]
    for cid in expired:
        del challenges[cid]

def cleanup_expired_sessions():
    """Remove expired sessions"""
    current_time = time.time()
    expired = [token for token, (_, _, exp) in sessions.items() if exp < current_time]
    for token in expired:
        del sessions[token]

@app.get("/ai", response_class=HTMLResponse)
async def serve_ai():
    with open("static/index.html") as f:
        return f.read()
    
@app.get("/", response_class=HTMLResponse)
async def serve_root():
    with open("static/index.html") as f:
        return f.read()

def hash_key(chat_key: str) -> str:
    return hashlib.sha256(chat_key.encode('utf-8')).hexdigest()

def hash_key_with_challenge(key_hash: str, challenge: str) -> str:
    """Create a challenge-response hash"""
    combined = f"{key_hash}:{challenge}"
    return hashlib.sha256(combined.encode('utf-8')).hexdigest()

@app.post("/challenge")
async def get_challenge(request: Request):
    """Generate a challenge for authentication"""
    cleanup_expired_challenges()
    
    data = await request.json()
    chat_id = data.get("chat_id")
    
    if not chat_id:
        raise HTTPException(status_code=400, detail="chat_id required")
    
    # Generate a random challenge
    challenge_id = secrets.token_urlsafe(32)
    challenge = secrets.token_urlsafe(32)
    
    # Store challenge with expiry
    challenges[challenge_id] = (challenge, time.time() + CHALLENGE_EXPIRY)
    
    return {
        "challenge_id": challenge_id,
        "challenge": challenge
    }

@app.post("/create")
async def create_chatroom(request: Request):
    """Create a new chatroom with hashed key"""
    data = await request.json()
    chat_id = data.get("chat_id")
    key_hash = data.get("key_hash")  # Client sends pre-hashed key
    
    if not chat_id or not key_hash:
        raise HTTPException(status_code=400, detail="chat_id and key_hash required")
    
    async with app.state.db.acquire() as conn:
        # Check if room already exists
        existing = await conn.fetchrow("SELECT chat_id FROM chatroom WHERE chat_id=$1", chat_id)
        if existing:
            raise HTTPException(status_code=409, detail="Chatroom already exists")
        
        # Create new room
        await conn.execute(
            "INSERT INTO chatroom (chat_id, key_hash) VALUES ($1, $2)",
            chat_id, key_hash
        )
        
    return {"status": "ok", "message": "Chatroom created successfully"}

@app.post("/enter")
async def enter_chatroom(request: Request):
    cleanup_expired_challenges()
    cleanup_expired_sessions()
    
    data = await request.json()
    chat_id = data.get("chat_id")
    username = data.get("username")
    challenge_id = data.get("challenge_id")
    response_hash = data.get("response")
    
    if not chat_id or not challenge_id or not response_hash or not username:
        raise HTTPException(status_code=400, detail="chat_id, username, challenge_id and response required")
    
    # Verify challenge exists and not expired
    if challenge_id not in challenges:
        raise HTTPException(status_code=401, detail="Invalid or expired challenge")
    
    challenge, expiry = challenges[challenge_id]
    if time.time() > expiry:
        del challenges[challenge_id]
        raise HTTPException(status_code=401, detail="Challenge expired")
    
    # Get stored key hash from database
    async with app.state.db.acquire() as conn:
        row = await conn.fetchrow("SELECT key_hash FROM chatroom WHERE chat_id=$1", chat_id)
        
        if row:
            # Verify the challenge-response
            expected_response = hash_key_with_challenge(row["key_hash"], challenge)
            
            if response_hash != expected_response:
                del challenges[challenge_id]
                raise HTTPException(status_code=403, detail="Wrong key for this chatroom")
            
            # Challenge used, delete it
            del challenges[challenge_id]
            
            # Create session token
            session_token = secrets.token_urlsafe(32)
            sessions[session_token] = (chat_id, username, time.time() + SESSION_EXPIRY)
            
            return {
                "status": "ok",
                "message": "Entered chatroom",
                "session_token": session_token
            }
        else:
            raise HTTPException(status_code=404, detail="Room not found")

from pydantic import BaseModel

class MessagesRequest(BaseModel):
    chat_id: str
    session_token: str
    before_id: int
    limit: int = 50

@app.post("/messages")
async def get_older_messages(request: MessagesRequest):
    cleanup_expired_sessions()
    
    chat_id = request.chat_id
    session_token = request.session_token
    before_id = request.before_id
    limit = min(request.limit, 100)

    # Verify session token
    if session_token not in sessions:
        raise HTTPException(status_code=401, detail="Invalid or expired session")
    
    stored_chat_id, stored_username, expiry = sessions[session_token]
    
    if time.time() > expiry:
        del sessions[session_token]
        raise HTTPException(status_code=401, detail="Session expired")
    
    if stored_chat_id != chat_id:
        raise HTTPException(status_code=403, detail="Session not valid for this chatroom")

    async with app.state.db.acquire() as conn:
        row = await conn.fetchrow("SELECT key_hash FROM chatroom WHERE chat_id=$1", chat_id)
        if not row:
            raise HTTPException(status_code=404, detail="Chatroom not found")

        try:
            rows = await conn.fetch(
                "SELECT id, username, message, sent_at, reply_to FROM chat_message WHERE chat_id=$1 AND id < $2 ORDER BY id DESC LIMIT $3",
                chat_id, before_id, limit
            )
        except Exception as e:
            logger.error(f"Error fetching older messages: {e}")
            raise HTTPException(status_code=500, detail="Failed to fetch messages")

        return [{
            "type": 'message',
            "message_id": row["id"],
            "username": row["username"],
            "msg": row["message"],
            "sent_at": (row["sent_at"] + timedelta(hours=3, minutes=30)).isoformat(),
            "reply_to": row["reply_to"]
        } for row in rows]

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}

    def add(self, chat_id: str, username: str, websocket: WebSocket):
        self.active_connections.setdefault(chat_id, {})[username] = websocket

    def remove(self, chat_id: str, username: str):
        if chat_id in self.active_connections:
            self.active_connections[chat_id].pop(username, None)
            if not self.active_connections[chat_id]:
                del self.active_connections[chat_id]

    def get_online_data(self, chat_id: str) -> dict:
        users = list(self.active_connections.get(chat_id, {}).keys())
        return {
            "type": "online",
            "users": users,
            "count": len(users),
        }

    async def broadcast(self, chat_id: str, message: dict):
        connections = self.active_connections.get(chat_id, {})
        dead_users = []

        for username, ws in connections.items():
            try:
                await ws.send_json(message)
            except Exception as e:
                logger.warning(f"Broadcast failed to {username}: {e}")
                dead_users.append(username)

        # cleanup after iteration
        for username in dead_users:
            self.remove(chat_id, username)

manager = ConnectionManager()

@app.websocket("/ws/{chat_id}/{username}/{session_token}")
async def websocket_endpoint(websocket: WebSocket, chat_id: str, username: str, session_token: str):
    # Verify session token before accepting connection
    cleanup_expired_sessions()
    
    if session_token not in sessions:
        await websocket.close(code=1008, reason="Invalid or expired session")
        return
    
    stored_chat_id, stored_username, expiry = sessions[session_token]
    
    if time.time() > expiry:
        del sessions[session_token]
        await websocket.close(code=1008, reason="Session expired")
        return
    
    if stored_chat_id != chat_id or stored_username != username:
        await websocket.close(code=1008, reason="Session not valid for this user/chatroom")
        return
    
    await websocket.accept()
    manager.add(chat_id, username, websocket)

    # notify online users
    await manager.broadcast(chat_id, manager.get_online_data(chat_id))

    try:
        # send recent messages
        async with app.state.db.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT id, username, message, sent_at, reply_to
                FROM chat_message
                WHERE chat_id=$1
                ORDER BY sent_at DESC
                LIMIT 50
                """,
                chat_id,
            )

        bulk = {
            "type": "bulk",
            "messages": [
                {
                    "type": "message",
                    "username": row["username"],
                    "msg": row["message"],
                    "message_id": row["id"],
                    "reply_to": row["reply_to"],
                    "sent_at": (row["sent_at"] + timedelta(hours=3, minutes=30)).isoformat(),
                }
                for row in reversed(rows)
            ],
        }

        await websocket.send_json(bulk)

        # main receive loop
        while True:
            data = await websocket.receive_json()

            msg = data.get("msg")
            reply_to = data.get("replyTo")
            local_id = data.get("localId")

            if not msg:
                logger.warning(f"Received empty message from {username} in chat_id={chat_id}")
                continue

            async with app.state.db.acquire() as conn:
                row = await conn.fetchrow(
                    """
                    INSERT INTO chat_message (chat_id, username, message, reply_to)
                    VALUES ($1, $2, $3, $4)
                    RETURNING id, sent_at
                    """,
                    chat_id,
                    username,
                    msg,
                    reply_to,
                )

            message_data = {
                "type": "message",
                "username": username,
                "msg": msg,
                "message_id": row['id'],
                "reply_to": reply_to,
                "sent_at": (row['sent_at'] + timedelta(hours=3, minutes=30)).isoformat(),
                "local_id": local_id,
            }

            logger.info(f"Broadcasting message {row['id']} in chat_id={chat_id}")
            await manager.broadcast(chat_id, message_data)

    except WebSocketDisconnect:
        logger.info(f"{username} disconnected from chat {chat_id}")

    except Exception as e:
        logger.error(f"WebSocket error for {username} in {chat_id}: {e}")

    finally:
        manager.remove(chat_id, username)
        await manager.broadcast(chat_id, manager.get_online_data(chat_id))