#!/usr/bin/env python3
"""
FastAPI application with multiple endpoints that publish to NATS JetStream
Each endpoint corresponds to an external API that workers will forward to
"""
import os
import json
import asyncio
from fastapi import FastAPI, HTTPException, Request, Body
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request as StarletteRequest
from starlette.responses import Response as StarletteResponse
from pydantic import BaseModel
import nats
from prometheus_client import Counter, Histogram, generate_latest, REGISTRY
from starlette.responses import Response
import uvicorn
from typing import Optional, Dict, Any, List, Union

# Prometheus metrics
request_count = Counter('http_requests_total', 'Total HTTP requests', ['method', 'endpoint', 'status'])
request_duration = Histogram('http_request_duration_seconds', 'HTTP request duration', ['method', 'endpoint'])

app = FastAPI(title="NATS Backend API - Multi-Endpoint", version="1.0.0")

# Custom CORS middleware to ensure headers are always added
class CORSHeaderMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: StarletteRequest, call_next):
        origin = request.headers.get("origin", "*")
        response = await call_next(request)
        response.headers["Access-Control-Allow-Origin"] = origin
        response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, PATCH, DELETE, OPTIONS"
        response.headers["Access-Control-Allow-Headers"] = "*"
        response.headers["Access-Control-Allow-Credentials"] = "false"
        response.headers["Access-Control-Max-Age"] = "600"
        return response

# Add custom CORS middleware first
app.add_middleware(CORSHeaderMiddleware)

# Also add FastAPI's CORS middleware as backup
app.add_middleware(
    CORSMiddleware,
    allow_origin_regex=r".*",  # Match all origins
    allow_credentials=False,
    allow_methods=["*"],  # Allow all HTTP methods
    allow_headers=["*"],  # Allow all headers
    expose_headers=["*"],  # Expose all headers
    max_age=600,
)

# NATS connection (will be initialized on startup)
nc = None
js = None

# Configuration
NATS_URL = os.getenv("NATS_URL", "nats://nats-server:4222")
NATS_USER = os.getenv("NATS_USER", "admin")
NATS_PASSWORD = os.getenv("NATS_PASSWORD", "package@321#")

# Endpoint to NATS subject mapping
ENDPOINT_ROUTES = {
    "/live/api/v1/deliveries/createdeliveries": "api.v1.deliveries.createdeliveries",
    "/live/api/v1/deliveries/updatedelivery": "api.v1.deliveries.updatedelivery",
    "/live/api/v2/partners/createriderlog": "api.v2.partners.createriderlog",
    "/live/api/v2/deliveries/createdeliverylog": "api.v2.deliveries.createdeliverylog",
    "/live/api/v2/partners/createbreaklog": "api.v2.partners.createbreaklog",
    "/live/api/v2/partners/updatebreaklog": "api.v2.partners.updatebreaklog",
    "/live/api/v1/mob/orders/createorder": "api.v1.mob.orders.createorder",
    "/live/api/v1/web/products/create": "api.v1.web.products.create",
    "/live/api/v1/mob/customers/create": "api.v1.mob.customers.create",
    "/live/api/v1/mob/customers/login": "api.v1.mob.customers.login",
}

@app.on_event("startup")
async def startup():
    """Initialize NATS connection on startup"""
    global nc, js
    try:
        print(f"Connecting to NATS at {NATS_URL}...")
        nc = await nats.connect(
            servers=[NATS_URL],
            user=NATS_USER,
            password=NATS_PASSWORD,
            reconnect_time_wait=2,
            max_reconnect_attempts=10
        )
        js = nc.jetstream()
        print("✅ Connected to NATS JetStream")
        print(f"✅ Configured {len(ENDPOINT_ROUTES)} endpoint routes")
    except Exception as e:
        print(f"❌ Failed to connect to NATS: {e}")
        raise

@app.on_event("shutdown")
async def shutdown():
    """Close NATS connection on shutdown"""
    global nc
    if nc:
        await nc.close()
        print("NATS connection closed")

@app.options("/{full_path:path}")
async def options_handler(full_path: str, request: Request):
    """Handle OPTIONS requests for CORS preflight"""
    origin = request.headers.get("origin")
    return JSONResponse(
        status_code=200,
        content={},
        headers={
            "Access-Control-Allow-Origin": origin if origin else "*",
            "Access-Control-Allow-Methods": "GET, POST, PUT, PATCH, DELETE, OPTIONS",
            "Access-Control-Allow-Headers": "*",
            "Access-Control-Allow-Credentials": "false",
            "Access-Control-Max-Age": "600",
        }
    )

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy", "nats_connected": nc.is_connected if nc else False}

@app.get("/ready")
async def ready():
    """Readiness check endpoint"""
    if nc and nc.is_connected:
        return {"status": "ready"}
    raise HTTPException(status_code=503, detail="Not ready")

async def publish_to_nats(endpoint: str, data: Union[Dict[str, Any], List[Dict[str, Any]]], request_method: str = "POST"):
    """Publish message to NATS with endpoint metadata"""
    payload = {
        "endpoint": endpoint,
        "method": request_method,
        "data": data,
        "received_at": int(asyncio.get_event_loop().time() * 1000),
        "original_path": endpoint
    }
    
    # Get NATS subject for this endpoint
    subject = ENDPOINT_ROUTES.get(endpoint, "api.unknown")
    
    if js:
        try:
            ack = await js.publish(subject, json.dumps(payload).encode())
            return ack.seq
        except Exception as e:
            print(f"❌ Failed to publish to NATS: {e}")
            raise HTTPException(status_code=500, detail=f"Failed to publish message: {str(e)}")
    else:
        raise HTTPException(status_code=503, detail="NATS not connected")

async def publish_to_nats_sync(endpoint: str, data: Union[Dict[str, Any], List[Dict[str, Any]]], request_method: str = "POST", timeout: int = 15):
    """
    Publish message to NATS and wait for a dedicated reply from the worker (Synchronous)
    """
    if not nc or not js:
        raise HTTPException(status_code=503, detail="NATS not connected")

    # Create a unique inbox for the reply
    reply_subject = nc.new_inbox()
    
    # Subscribe to the inbox matching the unique reply subject
    sub = await nc.subscribe(reply_subject)
    
    payload = {
        "endpoint": endpoint,
        "method": request_method,
        "data": data,
        "received_at": int(asyncio.get_event_loop().time() * 1000),
        "original_path": endpoint,
        "reply_subject": reply_subject  # Tell worker where to reply
    }
    
    # Get NATS subject for this endpoint
    subject = ENDPOINT_ROUTES.get(endpoint, "api.unknown")

    try:
        # Publish to the Stream
        ack = await js.publish(subject, json.dumps(payload).encode())
        print(f"Sync message published to {subject}, waiting for reply on {reply_subject}...")
        
        # Wait for the worker to process and reply
        # We assume the worker will send a single JSON message: {"status": <int>, "body": <str|json>}
        msg = await sub.next_msg(timeout=timeout)
        
        # Parse the reply
        response_data = json.loads(msg.data.decode())
        status_code = response_data.get("status", 500)
        body_content = response_data.get("body", "")
        
        # Attempt to parse body as JSON if it's a string, to avoid double-serialization
        # If the body is already a dict, use it directly.
        final_content = body_content
        if isinstance(body_content, str):
            try:
                final_content = json.loads(body_content)
            except:
                pass # Keep as string if not valid JSON

        await sub.unsubscribe()
        return status_code, final_content

    except asyncio.TimeoutError:
        await sub.unsubscribe()
        print(f"❌ Timeout waiting for sync response on {endpoint}")
        raise HTTPException(status_code=504, detail="Gateway Timeout: Worker took too long to respond")
    except Exception as e:
        await sub.unsubscribe()
        print(f"❌ Failed during sync processing: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to process synchronous request: {str(e)}")

# Endpoint 1: Update Delivery (v1) - External API requires PUT
@app.put("/live/api/v1/deliveries/updatedelivery")
async def update_delivery_v1(data: Dict[str, Any], request: Request):
    """Update Delivery endpoint - forwards to NATS (PUT only, as external API requires PUT)"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/deliveries/updatedelivery"
    method = "PUT"  # Always use PUT for this endpoint
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method="PUT", endpoint=endpoint).observe(duration)
        request_count.labels(method="PUT", endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method="PUT", endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint: Create Deliveries (v1)
@app.post("/live/api/v1/deliveries/createdeliveries")
async def create_deliveries_v1(data: Union[Dict[str, Any], List[Dict[str, Any]]], request: Request):
    """Create Deliveries endpoint - forwards payload to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/deliveries/createdeliveries"
    method = request.method
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 2: Create Rider Log (v2)
@app.post("/live/api/v2/partners/createriderlog")
async def create_rider_log_v2(data: Dict[str, Any], request: Request):
    """Create Rider Log endpoint - forwards to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v2/partners/createriderlog"
    method = request.method
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 3: Create Delivery Log (v2)
@app.post("/live/api/v2/deliveries/createdeliverylog")
async def create_delivery_log_v2(
    request: Request,
    data: List[Dict[str, Any]] = Body(
        ...,
        example=[{
            "logid": 0,
            "tenantid": 1,
            "partnerid": 44,
            "locationid": 1,
            "orderheaderid": 123456,
            "deliveryid": 654321,
            "userid": 1111,
            "orderid": "1-20231624",
            "orderstatus": "active",
            "starttime": "2025-12-10 17:51:03",
            "logdate": "2025-12-10 18:14:04",
            "latitude": "11.0050664",
            "longitude": "76.9508776"
        }]
    )
):
    """
    Create Delivery Log endpoint - forwards to NATS.
    Accepts either a single dict or a list of dicts to align with external API expectations.
    """
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v2/deliveries/createdeliverylog"
    method = request.method
    
    try:
        # Expect only a list of dicts
        if not isinstance(data, list) or not all(isinstance(item, dict) for item in data):
            raise HTTPException(status_code=422, detail="Body must be a list of objects")

        normalized_data = data

        message_id = await publish_to_nats(endpoint, normalized_data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 4: Create Break Rider Log (v2)
@app.post("/live/api/v2/partners/createbreaklog")
async def create_break_log_v2(data: Dict[str, Any], request: Request):
    """Create Break Rider Log endpoint - forwards to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v2/partners/createbreaklog"
    method = request.method
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 5: Update Break Rider Log (v2) - Supports both POST and PUT
@app.post("/live/api/v2/partners/updatebreaklog")
@app.put("/live/api/v2/partners/updatebreaklog")
async def update_break_log_v2(data: Dict[str, Any], request: Request):
    """Update Break Rider Log endpoint - forwards to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v2/partners/updatebreaklog"
    method = request.method  # Will be POST or PUT
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/metrics")
async def metrics():
    """Prometheus metrics endpoint"""
    return Response(content=generate_latest(REGISTRY), media_type="text/plain")

@app.get("/routes")
async def list_routes():
    """List all configured routes"""
    return {
        "routes": ENDPOINT_ROUTES,
        "total": len(ENDPOINT_ROUTES)
    }


# Endpoint 6: Create Order (Mob V1)
@app.post("/live/api/v1/mob/orders/createorder")
async def create_order_mob_v1(data: Dict[str, Any], request: Request):
    """Create Order endpoint (Mobile) - forwards to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/mob/orders/createorder"
    method = request.method
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 7: Create Product (Web V1)
@app.post("/live/api/v1/web/products/create")
async def create_product_web_v1(data: Dict[str, Any], request: Request):
    """Create Product endpoint (Web) - forwards to NATS"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/web/products/create"
    method = request.method
    
    try:
        message_id = await publish_to_nats(endpoint, data, method)
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status="200").inc()
        
        return JSONResponse(
            status_code=200,
            content={"status": "accepted", "message_id": message_id, "endpoint": endpoint}
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 8: Create Customer (Mob V1)
@app.post("/live/api/v1/mob/customers/create")
async def create_customer_v1(data: Dict[str, Any], request: Request):
    """Create Customer endpoint (Mobile) - synchronous wait for NATS worker"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/mob/customers/create"
    method = request.method
    
    try:
        status_code, content = await publish_to_nats_sync(endpoint, data, method)
        
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status=str(status_code)).inc()
        
        return JSONResponse(
            status_code=status_code,
            content=content
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint 9: Login Customer (Mob V1)
@app.post("/live/api/v1/mob/customers/login")
async def login_customer_v1(data: Dict[str, Any], request: Request):
    """Login Customer endpoint (Mobile) - synchronous wait for NATS worker"""
    start_time = asyncio.get_event_loop().time()
    endpoint = "/live/api/v1/mob/customers/login"
    method = request.method
    
    try:
        status_code, content = await publish_to_nats_sync(endpoint, data, method)
        
        duration = asyncio.get_event_loop().time() - start_time
        request_duration.labels(method=method, endpoint=endpoint).observe(duration)
        request_count.labels(method=method, endpoint=endpoint, status=str(status_code)).inc()
        
        return JSONResponse(
            status_code=status_code,
            content=content
        )
    except HTTPException:
        raise
    except Exception as e:
        request_count.labels(method=method, endpoint=endpoint, status="500").inc()
        print(f"❌ Error processing {endpoint}: {e}")
        raise HTTPException(status_code=500, detail=str(e))

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