"""
API Gateway for Energy Management Microservices
Central entry point that routes requests to appropriate microservices.
Port: 8000
"""

import asyncio
import aiohttp
from datetime import datetime
from fastapi import FastAPI, HTTPException, WebSocket, Depends, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
import logging
import json
from typing import Dict, Any, Optional
import os

from models import ServiceConfig, HealthResponse, GatewayStats
from service_registry import ServiceRegistry
from load_balancer import LoadBalancer
from auth_middleware import AuthMiddleware

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager"""
    logger.info("API Gateway starting up...")

    # Initialize service registry
    await service_registry.initialize()

    # Start health check task
    asyncio.create_task(health_check_task())

    logger.info("API Gateway startup complete")

    yield

    logger.info("API Gateway shutting down...")
    await service_registry.close()
    logger.info("API Gateway shutdown complete")

app = FastAPI(
    title="Energy Management API Gateway",
    description="Central API gateway for energy management microservices",
    version="1.0.0",
    lifespan=lifespan
)

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

# Service registry and load balancer
service_registry = ServiceRegistry()
load_balancer = LoadBalancer()
auth_middleware = AuthMiddleware()

# Service configuration
SERVICES = {
    "token-service": ServiceConfig(
        name="token-service",
        base_url=os.getenv("TOKEN_SERVICE_URL", "http://energy-token-service:8001"),
        health_endpoint="/health",
        auth_required=False
    ),
    "battery-service": ServiceConfig(
        name="battery-service",
        base_url=os.getenv("BATTERY_SERVICE_URL", "http://energy-battery-service:8002"),
        health_endpoint="/health",
        auth_required=True
    ),
    "demand-response-service": ServiceConfig(
        name="demand-response-service",
        base_url=os.getenv("DEMAND_RESPONSE_SERVICE_URL", "http://energy-demand-response-service:8003"),
        health_endpoint="/health",
        auth_required=True
    ),
    "p2p-trading-service": ServiceConfig(
        name="p2p-trading-service",
        base_url=os.getenv("P2P_TRADING_SERVICE_URL", "http://energy-p2p-trading-service:8004"),
        health_endpoint="/health",
        auth_required=True
    ),
    "forecasting-service": ServiceConfig(
        name="forecasting-service",
        base_url=os.getenv("FORECASTING_SERVICE_URL", "http://energy-forecasting-service:8005"),
        health_endpoint="/health",
        auth_required=True
    ),
    "iot-control-service": ServiceConfig(
        name="iot-control-service",
        base_url=os.getenv("IOT_CONTROL_SERVICE_URL", "http://energy-iot-control-service:8006"),
        health_endpoint="/health",
        auth_required=True
    ),
    "sensor-service": ServiceConfig(
        name="sensor-service",
        base_url=os.getenv("SENSOR_SERVICE_URL", "http://energy-sensor-service:8007"),
        health_endpoint="/health",
        auth_required=True
    ),
    "data-ingestion-service": ServiceConfig(
        name="data-ingestion-service",
        base_url=os.getenv("DATA_INGESTION_SERVICE_URL", "http://energy-data-ingestion-service:8008"),
        health_endpoint="/health",
        auth_required=False
    )
}

# Request statistics
request_stats = {
    "total_requests": 0,
    "successful_requests": 0,
    "failed_requests": 0,
    "service_requests": {service: 0 for service in SERVICES.keys()},
    "start_time": datetime.utcnow()
}

@app.get("/health", response_model=HealthResponse)
async def gateway_health_check():
    """Gateway health check endpoint"""
    try:
        # Check all services
        service_health = await service_registry.get_all_service_health()

        healthy_services = sum(1 for status in service_health.values() if status.get("status") == "healthy")
        total_services = len(SERVICES)

        overall_status = "healthy" if healthy_services == total_services else "degraded"

        return HealthResponse(
            service="api-gateway",
            status=overall_status,
            timestamp=datetime.utcnow(),
            version="1.0.0",
            services=service_health,
            healthy_services=healthy_services,
            total_services=total_services
        )
    except Exception as e:
        logger.error(f"Gateway health check failed: {e}")
        raise HTTPException(status_code=503, detail="Service Unavailable")

@app.get("/services/status")
async def get_services_status():
    """Get status of all registered services"""
    try:
        service_health = await service_registry.get_all_service_health()
        return {
            "services": service_health,
            "timestamp": datetime.utcnow().isoformat(),
            "total_services": len(SERVICES),
            "healthy_services": sum(1 for status in service_health.values() if status.get("status") == "healthy")
        }
    except Exception as e:
        logger.error(f"Error getting services status: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.get("/stats", response_model=GatewayStats)
async def get_gateway_stats():
    """Get API gateway statistics"""
    uptime = (datetime.utcnow() - request_stats["start_time"]).total_seconds()

    return GatewayStats(
        total_requests=request_stats["total_requests"],
        successful_requests=request_stats["successful_requests"],
        failed_requests=request_stats["failed_requests"],
        success_rate=round((request_stats["successful_requests"] / max(request_stats["total_requests"], 1)) * 100, 2),
        uptime_seconds=uptime,
        service_requests=request_stats["service_requests"],
        timestamp=datetime.utcnow()
    )

# Token Service Routes
@app.api_route("/api/v1/tokens/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def token_service_proxy(request: Request, path: str):
    """Proxy requests to token service"""
    return await proxy_request(request, "token-service", f"/{path}")

# Battery Service Routes
@app.api_route("/api/v1/batteries/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def battery_service_proxy(request: Request, path: str):
    """Proxy requests to battery service"""
    return await proxy_request(request, "battery-service", f"/{path}")

# Demand Response Service Routes
@app.api_route("/api/v1/demand-response/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def demand_response_service_proxy(request: Request, path: str):
    """Proxy requests to demand response service"""
    return await proxy_request(request, "demand-response-service", f"/{path}")

# P2P Trading Service Routes
@app.api_route("/api/v1/p2p/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def p2p_trading_service_proxy(request: Request, path: str):
    """Proxy requests to P2P trading service"""
    return await proxy_request(request, "p2p-trading-service", f"/{path}")

# Forecasting Service Routes
@app.api_route("/api/v1/forecast/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def forecasting_service_proxy(request: Request, path: str):
    """Proxy requests to forecasting service"""
    return await proxy_request(request, "forecasting-service", f"/{path}")

# IoT Control Service Routes
@app.api_route("/api/v1/iot/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def iot_control_service_proxy(request: Request, path: str):
    """Proxy requests to IoT control service"""
    return await proxy_request(request, "iot-control-service", f"/{path}")

# Sensor Service Routes (Original Dashboard Functionality)
@app.api_route("/api/v1/sensors/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def sensor_service_proxy(request: Request, path: str):
    """Proxy requests to sensor service"""
    return await proxy_request(request, "sensor-service", f"/{path}")

# Data Ingestion Service Routes (SA4CPS FTP Monitoring)
@app.api_route("/api/v1/ingestion/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def data_ingestion_service_proxy(request: Request, path: str):
    """Proxy requests to data ingestion service"""
    return await proxy_request(request, "data-ingestion-service", f"/{path}")

@app.api_route("/api/v1/sources/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def data_sources_proxy(request: Request, path: str):
    """Proxy requests to data ingestion service for data sources"""
    return await proxy_request(request, "data-ingestion-service", f"/sources/{path}")

@app.get("/api/v1/sources")
async def data_sources_list_proxy(request: Request):
    """Proxy requests to data ingestion service for sources list"""
    return await proxy_request(request, "data-ingestion-service", "/sources")

@app.api_route("/api/v1/rooms/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def room_service_proxy(request: Request, path: str):
    """Proxy requests to sensor service for room management"""
    return await proxy_request(request, "sensor-service", f"/rooms/{path}")

@app.api_route("/api/v1/data/{path:path}", methods=["GET", "POST"])
async def data_service_proxy(request: Request, path: str):
    """Proxy requests to sensor service for data operations"""
    return await proxy_request(request, "sensor-service", f"/data/{path}")

@app.api_route("/api/v1/analytics/{path:path}", methods=["GET", "POST"])
async def analytics_service_proxy(request: Request, path: str):
    """Proxy requests to sensor service for analytics"""
    return await proxy_request(request, "sensor-service", f"/analytics/{path}")

@app.api_route("/api/v1/export", methods=["GET"])
async def export_service_proxy(request: Request):
    """Proxy requests to sensor service for data export"""
    return await proxy_request(request, "sensor-service", "/export")

@app.api_route("/api/v1/events", methods=["GET"])
async def events_service_proxy(request: Request):
    """Proxy requests to sensor service for system events"""
    return await proxy_request(request, "sensor-service", "/events")

# WebSocket proxy for real-time data
@app.websocket("/ws")
async def websocket_proxy(websocket: WebSocket):
    """Proxy WebSocket connections to sensor service"""
    try:
        # Get sensor service URL
        service_url = await load_balancer.get_service_url("sensor-service")
        if not service_url:
            await websocket.close(code=1003, reason="Sensor service unavailable")
            return

        # For simplicity, we'll just accept the connection and forward to sensor service
        # In a production setup, you'd want a proper WebSocket proxy
        await websocket.accept()

        # For now, we'll handle this by having the sensor service manage WebSockets directly
        # The frontend should connect to the sensor service WebSocket endpoint directly
        await websocket.send_text(json.dumps({
            "type": "proxy_info",
            "message": "Connect directly to sensor service WebSocket at /ws",
            "sensor_service_url": service_url.replace("http://", "ws://") + "/ws"
        }))

        # Keep connection alive
        while True:
            try:
                await websocket.receive_text()
            except:
                break

    except Exception as e:
        logger.error(f"WebSocket proxy error: {e}")

async def proxy_request(request: Request, service_name: str, path: str):
    """Generic request proxy function"""
    try:
        # Update request statistics
        request_stats["total_requests"] += 1
        request_stats["service_requests"][service_name] += 1

        # Get service configuration
        service_config = SERVICES.get(service_name)
        if not service_config:
            raise HTTPException(status_code=404, detail=f"Service {service_name} not found")

        # Check authentication if required
        if service_config.auth_required:
            await auth_middleware.verify_token(request)

        # Get healthy service instance
        service_url = await load_balancer.get_service_url(service_name)

        # Prepare request
        url = f"{service_url}{path}"
        method = request.method
        headers = dict(request.headers)

        # Remove hop-by-hop headers
        headers.pop("host", None)
        headers.pop("content-length", None)

        # Get request body
        body = None
        if method in ["POST", "PUT", "PATCH"]:
            body = await request.body()

        # Make request to service
        async with aiohttp.ClientSession() as session:
            async with session.request(
                method=method,
                url=url,
                headers=headers,
                data=body,
                params=dict(request.query_params),
                timeout=aiohttp.ClientTimeout(total=30)
            ) as response:

                # Get response data
                response_data = await response.read()
                response_headers = dict(response.headers)

                # Remove hop-by-hop headers from response
                response_headers.pop("transfer-encoding", None)
                response_headers.pop("connection", None)

                # Update success statistics
                if response.status < 400:
                    request_stats["successful_requests"] += 1
                else:
                    request_stats["failed_requests"] += 1

                # Return response
                return Response(
                    content=response_data,
                    status_code=response.status,
                    headers=response_headers,
                    media_type=response_headers.get("content-type")
                )

    except aiohttp.ClientError as e:
        request_stats["failed_requests"] += 1
        logger.error(f"Service {service_name} connection error: {e}")
        raise HTTPException(status_code=503, detail=f"Service {service_name} unavailable")

    except HTTPException:
        request_stats["failed_requests"] += 1
        raise

    except Exception as e:
        request_stats["failed_requests"] += 1
        logger.error(f"Proxy error for {service_name}: {e}")
        raise HTTPException(status_code=500, detail="Internal gateway error")

@app.get("/api/v1/overview")
async def get_system_overview():
    """Get comprehensive system overview from all services"""
    try:
        overview = {}

        # Get data from each service
        for service_name in SERVICES.keys():
            try:
                if await service_registry.is_service_healthy(service_name):
                    service_url = await load_balancer.get_service_url(service_name)

                    async with aiohttp.ClientSession() as session:
                        # Try to get service-specific overview data
                        overview_endpoints = {
                            "sensor-service": "/analytics/summary",
                            "battery-service": "/batteries",
                            "demand-response-service": "/flexibility/current",
                            "p2p-trading-service": "/market/status",
                            "forecasting-service": "/forecast/summary",
                            "iot-control-service": "/devices/summary"
                        }

                        endpoint = overview_endpoints.get(service_name)
                        if endpoint:
                            async with session.get(f"{service_url}{endpoint}", timeout=aiohttp.ClientTimeout(total=5)) as response:
                                if response.status == 200:
                                    data = await response.json()
                                    overview[service_name] = data
                                else:
                                    overview[service_name] = {"status": "error", "message": "Service returned error"}
                        else:
                            overview[service_name] = {"status": "available"}

            except Exception as e:
                logger.warning(f"Could not get overview from {service_name}: {e}")
                overview[service_name] = {"status": "unavailable", "error": str(e)}

        return {
            "system_overview": overview,
            "timestamp": datetime.utcnow().isoformat(),
            "services_checked": len(SERVICES)
        }

    except Exception as e:
        logger.error(f"Error getting system overview: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

async def health_check_task():
    """Background task for periodic health checks"""
    logger.info("Starting health check task")

    while True:
        try:
            await service_registry.update_all_service_health()
            await asyncio.sleep(30)  # Check every 30 seconds

        except Exception as e:
            logger.error(f"Error in health check task: {e}")
            await asyncio.sleep(60)

# Initialize service registry with services
asyncio.create_task(service_registry.register_services(SERVICES))

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