import json
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from app.api import auth, chats, messages, media, notifications, encryption
from app.db.session import create_tables
from app.websocket.connection_manager import connection_manager
from app.websocket.chat_handler import chat_handler

app = FastAPI(
    title="Real-time Chat API",
    description="A comprehensive real-time chat API with WebSocket support, media sharing, E2E encryption, and push notifications",
    version="1.0.0",
    openapi_url="/openapi.json"
)

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

# Create database tables
create_tables()

# Include API routers
app.include_router(auth.router, prefix="/api/auth", tags=["authentication"])
app.include_router(chats.router, prefix="/api/chats", tags=["chats"])
app.include_router(messages.router, prefix="/api/messages", tags=["messages"])
app.include_router(media.router, prefix="/api/media", tags=["media"])
app.include_router(notifications.router, prefix="/api/notifications", tags=["notifications"])
app.include_router(encryption.router, prefix="/api/encryption", tags=["encryption"])

@app.get("/")
def read_root():
    return {
        "title": "Real-time Chat API",
        "description": "A comprehensive real-time chat API with WebSocket support",
        "version": "1.0.0",
        "features": [
            "Real-time messaging with WebSocket",
            "Direct messages and group chats",
            "Media sharing (images, videos, documents)",
            "End-to-end encryption",
            "Push notifications",
            "Mention alerts",
            "Message read receipts",
            "Typing indicators"
        ],
        "documentation": "/docs",
        "websocket_endpoint": "/ws/{token}",
        "health_check": "/health"
    }

@app.get("/health")
def health_check():
    return {
        "status": "healthy",
        "service": "Real-time Chat API",
        "version": "1.0.0",
        "features": {
            "websocket": "enabled",
            "media_upload": "enabled",
            "encryption": "enabled",
            "notifications": "enabled"
        }
    }

@app.websocket("/ws/{token}")
async def websocket_endpoint(websocket: WebSocket, token: str):
    """WebSocket endpoint for real-time chat"""
    user_id = await connection_manager.connect(websocket, token)
    if not user_id:
        return
    
    try:
        while True:
            # Receive message from client
            data = await websocket.receive_text()
            
            try:
                message_data = json.loads(data)
                await chat_handler.handle_message(websocket, message_data)
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({
                    "type": "error",
                    "message": "Invalid JSON format"
                }))
            except Exception as e:
                await websocket.send_text(json.dumps({
                    "type": "error",
                    "message": f"Error processing message: {str(e)}"
                }))
                
    except WebSocketDisconnect:
        await connection_manager.disconnect(websocket)
    except Exception as e:
        print(f"WebSocket error: {e}")
        await connection_manager.disconnect(websocket)

@app.get("/api/status")
def get_system_status():
    """Get system status and statistics"""
    return {
        "active_connections": len(connection_manager.active_connections),
        "active_chats": len(connection_manager.chat_members),
        "server_time": "2024-12-20T12:00:00Z",
        "uptime": "0 days, 0 hours, 0 minutes"
    }

# Add startup event
@app.on_event("startup")
async def startup_event():
    print("🚀 Real-time Chat API starting up...")
    print("📡 WebSocket endpoint: /ws/{token}")
    print("📖 API Documentation: /docs")
    print("💚 Health check: /health")

# Add shutdown event  
@app.on_event("shutdown")
async def shutdown_event():
    print("🛑 Real-time Chat API shutting down...")
    # Clean up active connections
    for user_id, connections in connection_manager.active_connections.items():
        for websocket in connections:
            try:
                await websocket.close()
            except Exception:
                pass