"""
API Gateway - Central entry point for all microservices
Handles routing, authentication, rate limiting, and cross-cutting concerns
"""

import asyncio
import json
import structlog
from fastapi import FastAPI, Request, HTTPException, Depends, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
import httpx
import time
import redis.asyncio as aioredis
from typing import Dict, Any

from app.core.config import settings
from app.core.service_discovery import ServiceDiscovery
from app.middleware.auth import AuthMiddleware
from app.middleware.logging import LoggingMiddleware
from app.middleware.rate_limit import RateLimitMiddleware
from app.middleware.subscription import SubscriptionMiddleware
from app.middleware.demo_middleware import DemoMiddleware
from app.routes import auth, tenant, notification, nominatim, user, subscription, demo, pos
from shared.monitoring.logging import setup_logging
from shared.monitoring.metrics import MetricsCollector

# Setup logging
setup_logging("gateway", settings.LOG_LEVEL)
logger = structlog.get_logger()

# Create FastAPI app
app = FastAPI(
    title="Bakery Forecasting API Gateway",
    description="Central API Gateway for bakery forecasting microservices",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# Initialize metrics collector
metrics_collector = MetricsCollector("gateway")

# Service discovery
service_discovery = ServiceDiscovery()

# Redis client for SSE streaming
redis_client = None

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

# Custom middleware - Add in REVERSE order (last added = first executed)
# Execution order: DemoMiddleware -> AuthMiddleware -> SubscriptionMiddleware -> RateLimitMiddleware -> LoggingMiddleware
app.add_middleware(LoggingMiddleware)          # Executes 5th (outermost)
app.add_middleware(RateLimitMiddleware, calls_per_minute=300)  # Executes 4th
app.add_middleware(SubscriptionMiddleware, tenant_service_url=settings.TENANT_SERVICE_URL)  # Executes 3rd
app.add_middleware(AuthMiddleware)             # Executes 2nd - Checks for demo context
app.add_middleware(DemoMiddleware)             # Executes 1st (innermost) - Sets demo user context FIRST

# Include routers
app.include_router(auth.router, prefix="/api/v1/auth", tags=["authentication"])
app.include_router(user.router, prefix="/api/v1/users", tags=["users"])
app.include_router(tenant.router, prefix="/api/v1/tenants", tags=["tenants"])
app.include_router(subscription.router, prefix="/api/v1", tags=["subscriptions"])
app.include_router(notification.router, prefix="/api/v1/notifications", tags=["notifications"])
app.include_router(nominatim.router, prefix="/api/v1/nominatim", tags=["location"])
app.include_router(pos.router, prefix="/api/v1/pos", tags=["pos"])
app.include_router(demo.router, prefix="/api/v1", tags=["demo"])


@app.on_event("startup")
async def startup_event():
    """Application startup"""
    global redis_client
    
    logger.info("Starting API Gateway")
    
    # Connect to Redis for SSE streaming
    try:
        redis_client = aioredis.from_url(settings.REDIS_URL)
        logger.info("Connected to Redis for SSE streaming")
    except Exception as e:
        logger.error(f"Failed to connect to Redis: {e}")

    metrics_collector.register_counter(
        "gateway_auth_requests_total",
        "Total authentication requests"
    )
    metrics_collector.register_counter(
        "gateway_auth_responses_total", 
        "Total authentication responses"
    )
    metrics_collector.register_counter(
        "gateway_auth_errors_total",
        "Total authentication errors"
    )
    
    metrics_collector.register_histogram(
        "gateway_request_duration_seconds",
        "Request duration in seconds"
    )
    
    logger.info("Metrics registered successfully")
    
    metrics_collector.start_metrics_server(8080)
    
    logger.info("API Gateway started successfully")

@app.on_event("shutdown")
async def shutdown_event():
    """Application shutdown"""
    global redis_client
    
    logger.info("Shutting down API Gateway")
    
    # Close Redis connection
    if redis_client:
        await redis_client.close()
    
    # Clean up service discovery
    # await service_discovery.cleanup()
    
    logger.info("API Gateway shutdown complete")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "api-gateway",
        "version": "1.0.0",
        "timestamp": time.time()
    }

@app.get("/metrics")
async def metrics():
    """Metrics endpoint for monitoring"""
    return {"metrics": "enabled"}

# ================================================================
# SERVER-SENT EVENTS (SSE) ENDPOINT
# ================================================================

@app.get("/api/events")
async def events_stream(request: Request, tenant_id: str):
    """
    Server-Sent Events stream for real-time notifications.

    Authentication is handled by auth middleware via query param token.
    User context is available in request.state.user (injected by middleware).
    Tenant ID is provided by the frontend as a query parameter.
    """
    global redis_client

    if not redis_client:
        raise HTTPException(status_code=503, detail="SSE service unavailable")

    # Extract user context from request state (set by auth middleware)
    user_context = request.state.user
    user_id = user_context.get('user_id')
    email = user_context.get('email')

    # Validate tenant_id parameter
    if not tenant_id:
        raise HTTPException(status_code=400, detail="tenant_id query parameter is required")

    logger.info(f"SSE connection request for user {email}, tenant {tenant_id}")
    
    logger.info(f"SSE connection established for tenant: {tenant_id}")
    
    async def event_generator():
        """Generate server-sent events from Redis pub/sub"""
        pubsub = None
        try:
            # Subscribe to tenant-specific alert channel
            pubsub = redis_client.pubsub()
            channel_name = f"alerts:{tenant_id}"
            await pubsub.subscribe(channel_name)
            
            # Send initial connection event
            yield f"event: connection\n"
            yield f"data: {json.dumps({'type': 'connected', 'message': 'SSE connection established', 'timestamp': time.time()})}\n\n"
            
            heartbeat_counter = 0
            
            while True:
                # Check if client has disconnected
                if await request.is_disconnected():
                    logger.info(f"SSE client disconnected for tenant: {tenant_id}")
                    break
                
                try:
                    # Get message from Redis with timeout
                    message = await asyncio.wait_for(pubsub.get_message(ignore_subscribe_messages=True), timeout=10.0)
                    
                    if message and message['type'] == 'message':
                        # Forward the alert/notification from Redis
                        alert_data = json.loads(message['data'])

                        # Determine event type based on alert data
                        event_type = "notification"
                        if alert_data.get('item_type') == 'alert':
                            if alert_data.get('severity') in ['high', 'urgent']:
                                event_type = "inventory_alert"
                            else:
                                event_type = "notification"
                        elif alert_data.get('item_type') == 'recommendation':
                            event_type = "notification"

                        yield f"event: {event_type}\n"
                        yield f"data: {json.dumps(alert_data)}\n\n"
                        
                        logger.debug(f"SSE message sent to tenant {tenant_id}: {alert_data.get('title')}")
                    
                except asyncio.TimeoutError:
                    # Send heartbeat every 10 timeouts (100 seconds)
                    heartbeat_counter += 1
                    if heartbeat_counter >= 10:
                        yield f"event: heartbeat\n"
                        yield f"data: {json.dumps({'type': 'heartbeat', 'timestamp': time.time()})}\n\n"
                        heartbeat_counter = 0
                    
        except asyncio.CancelledError:
            logger.info(f"SSE connection cancelled for tenant: {tenant_id}")
        except Exception as e:
            logger.error(f"SSE error for tenant {tenant_id}: {e}")
        finally:
            if pubsub:
                await pubsub.unsubscribe()
                await pubsub.close()
            logger.info(f"SSE connection closed for tenant: {tenant_id}")
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control",
        }
    )

# ================================================================
# WEBSOCKET ROUTING FOR TRAINING SERVICE
# ================================================================

@app.websocket("/api/v1/tenants/{tenant_id}/training/jobs/{job_id}/live")
async def websocket_training_progress(websocket: WebSocket, tenant_id: str, job_id: str):
    """
    WebSocket proxy that forwards connections directly to training service.
    Acts as a pure proxy - does NOT handle websocket logic, just forwards to training service.
    All auth, message handling, and business logic is in the training service.
    """
    # Get token from query params (required for training service authentication)
    token = websocket.query_params.get("token")
    if not token:
        logger.warning(f"WebSocket proxy rejected - missing token for job {job_id}")
        await websocket.accept()
        await websocket.close(code=1008, reason="Authentication token required")
        return

    # Accept the connection immediately
    await websocket.accept()

    logger.info(f"Gateway proxying WebSocket to training service for job {job_id}, tenant {tenant_id}")

    # Build WebSocket URL to training service - forward to the exact same path
    training_service_base = settings.TRAINING_SERVICE_URL.rstrip('/')
    training_ws_url = training_service_base.replace('http://', 'ws://').replace('https://', 'wss://')
    training_ws_url = f"{training_ws_url}/api/v1/tenants/{tenant_id}/training/jobs/{job_id}/live?token={token}"

    training_ws = None

    try:
        # Connect to training service WebSocket
        import websockets

        training_ws = await websockets.connect(
            training_ws_url,
            ping_interval=None,    # Let training service handle heartbeat
            ping_timeout=None,
            close_timeout=10,
            open_timeout=30,       # Allow time for training service to setup
            max_size=2**20,
            max_queue=32
        )

        logger.info(f"Gateway connected to training service WebSocket for job {job_id}")

        async def forward_to_training():
            """Forward messages from frontend to training service"""
            try:
                while training_ws and training_ws.open:
                    data = await websocket.receive()

                    if data.get("type") == "websocket.receive":
                        if "text" in data:
                            await training_ws.send(data["text"])
                            logger.debug(f"Gateway forwarded frontend->training: {data['text'][:100]}")
                        elif "bytes" in data:
                            await training_ws.send(data["bytes"])
                    elif data.get("type") == "websocket.disconnect":
                        logger.info(f"Frontend disconnected for job {job_id}")
                        break
            except Exception as e:
                logger.error(f"Error forwarding frontend->training for job {job_id}: {e}")

        async def forward_to_frontend():
            """Forward messages from training service to frontend"""
            try:
                while training_ws and training_ws.open:
                    message = await training_ws.recv()
                    await websocket.send_text(message)
                    logger.debug(f"Gateway forwarded training->frontend: {message[:100]}")
            except Exception as e:
                logger.error(f"Error forwarding training->frontend for job {job_id}: {e}")

        # Run both forwarding tasks concurrently
        await asyncio.gather(
            forward_to_training(),
            forward_to_frontend(),
            return_exceptions=True
        )

    except websockets.exceptions.ConnectionClosedError as e:
        logger.warning(f"Training service WebSocket closed for job {job_id}: {e}")
    except websockets.exceptions.WebSocketException as e:
        logger.error(f"WebSocket exception for job {job_id}: {e}")
    except Exception as e:
        logger.error(f"WebSocket proxy error for job {job_id}: {e}")
    finally:
        # Cleanup
        if training_ws and not training_ws.closed:
            try:
                await training_ws.close()
                logger.info(f"Closed training service WebSocket for job {job_id}")
            except Exception as e:
                logger.warning(f"Error closing training service WebSocket for job {job_id}: {e}")

        try:
            if not websocket.client_state.name == 'DISCONNECTED':
                await websocket.close(code=1000, reason="Proxy closed")
        except Exception as e:
            logger.warning(f"Error closing frontend WebSocket for job {job_id}: {e}")

        logger.info(f"Gateway WebSocket proxy cleanup completed for job {job_id}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)