#!/usr/bin/env python3
"""
NATS Worker that consumes messages and routes them to different external endpoints
Based on the endpoint in the message, forwards to the corresponding external API
"""
import os
import json
import asyncio
import aiohttp
import nats
from prometheus_client import Counter, Histogram, Gauge, start_http_server
import signal
import sys

# Prometheus metrics
messages_processed = Counter('worker_messages_processed_total', 'Total messages processed', ['status', 'endpoint'])
message_duration = Histogram('worker_message_duration_seconds', 'Message processing duration', ['endpoint'])
messages_in_flight = Gauge('worker_messages_in_flight', 'Messages currently being processed')
queue_depth = Gauge('worker_queue_depth', 'Approximate queue depth')

# 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#")
NATS_STREAM = os.getenv("NATS_STREAM", "EVENTS")
NATS_SUBJECT = os.getenv("NATS_SUBJECT", "api.>")  # Subscribe to all API subjects
NATS_CONSUMER = os.getenv("NATS_CONSUMER", "worker_consumer")
WORKER_CONCURRENCY = int(os.getenv("WORKER_CONCURRENCY", "10"))
RETRY_ATTEMPTS = int(os.getenv("RETRY_ATTEMPTS", "5"))
RETRY_DELAY = int(os.getenv("RETRY_DELAY_SECONDS", "5"))

# Base URL for external APIs
BASE_URL = os.getenv("EXTERNAL_BASE_URL", "https://jupiter.nearle.app")

# Endpoint mapping: FastAPI endpoint -> External API endpoint
ENDPOINT_MAPPING = {
    "/live/api/v1/deliveries/createdeliveries": f"{BASE_URL}/live/api/v1/deliveries/createdeliveries",
    "/live/api/v1/deliveries/updatedelivery": f"{BASE_URL}/live/api/v1/deliveries/updatedelivery",
    "/live/api/v2/partners/createriderlog": f"{BASE_URL}/live/api/v2/partners/createriderlog",
    "/live/api/v2/deliveries/createdeliverylog": f"{BASE_URL}/live/api/v2/deliveries/createdeliverylog",
    "/live/api/v2/partners/createbreaklog": f"{BASE_URL}/live/api/v2/partners/createbreaklog",
    "/live/api/v2/partners/updatebreaklog": f"{BASE_URL}/live/api/v2/partners/updatebreaklog",
}

# Global state
nc = None
js = None
running = True
semaphore = asyncio.Semaphore(WORKER_CONCURRENCY)

def signal_handler(sig, frame):
    """Handle shutdown signals"""
    global running
    print("\n🛑 Shutdown signal received, stopping worker...")
    running = False

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

async def forward_to_external(endpoint: str, payload: dict, api_key: str = None, method: str = "POST") -> bool:
    """Forward message to external endpoint with retries"""
    external_url = ENDPOINT_MAPPING.get(endpoint)
    
    if not external_url:
        print(f"⚠️  Unknown endpoint: {endpoint}, skipping...")
        return False
    
    headers = {
        "Content-Type": "application/json"
    }
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    
    # Get the actual data to forward
    data_to_forward = payload.get("data", payload)

    # Normalize payload shape for endpoints that expect an array
    if endpoint == "/live/api/v2/deliveries/createdeliverylog":
        if isinstance(data_to_forward, dict):
            data_to_forward = [data_to_forward]
        elif not isinstance(data_to_forward, list):
            # Fallback to wrapping anything that's not already list-like
            data_to_forward = [data_to_forward]

        # Debug logging for delivery log payload shape
        print(f"➡️  Forwarding createdeliverylog payload: {data_to_forward}")
    
    # Normalize method to uppercase
    http_method = method.upper() if method else "POST"
    
    for attempt in range(RETRY_ATTEMPTS):
        try:
            async with aiohttp.ClientSession() as session:
                # Use the correct HTTP method
                if http_method == "GET":
                    async with session.get(
                        external_url,
                        params=data_to_forward if isinstance(data_to_forward, dict) else {},
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30)
                    ) as response:
                        return await _handle_response(response, endpoint, attempt)
                elif http_method == "PUT":
                    async with session.put(
                        external_url,
                        json=data_to_forward,
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30)
                    ) as response:
                        return await _handle_response(response, endpoint, attempt)
                elif http_method == "PATCH":
                    async with session.patch(
                        external_url,
                        json=data_to_forward,
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30)
                    ) as response:
                        return await _handle_response(response, endpoint, attempt)
                elif http_method == "DELETE":
                    async with session.delete(
                        external_url,
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30)
                    ) as response:
                        return await _handle_response(response, endpoint, attempt)
                else:  # Default to POST
                    async with session.post(
                        external_url,
                        json=data_to_forward,
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30)
                    ) as response:
                        return await _handle_response(response, endpoint, attempt)
        except asyncio.TimeoutError:
            print(f"⚠️  Timeout forwarding {endpoint} to external API (attempt {attempt + 1}/{RETRY_ATTEMPTS})")
            if attempt < RETRY_ATTEMPTS - 1:
                await asyncio.sleep(RETRY_DELAY * (attempt + 1))
        except Exception as e:
            print(f"⚠️  Error forwarding {endpoint} to external API: {e} (attempt {attempt + 1}/{RETRY_ATTEMPTS})")
            if attempt < RETRY_ATTEMPTS - 1:
                await asyncio.sleep(RETRY_DELAY * (attempt + 1))
    
    return False

async def _handle_response(response, endpoint: str, attempt: int) -> bool:
    """Handle HTTP response and determine if we should retry"""
    body_text = None
    try:
        body_text = await response.text()
    except Exception:
        body_text = "<no body>"

    if response.status == 200 or response.status == 201:
        return True
    elif response.status >= 500:
        # Server error - retry
        print(f"⚠️  External API returned {response.status} for {endpoint}, retrying... (attempt {attempt + 1}/{RETRY_ATTEMPTS}) body={body_text}")
        if attempt < RETRY_ATTEMPTS - 1:
            await asyncio.sleep(RETRY_DELAY * (attempt + 1))  # Exponential backoff
        return False
    else:
        # Client error - don't retry
        print(f"❌ External API returned {response.status} for {endpoint}, not retrying; body={body_text}")
        return False

async def process_message(msg):
    """Process a single message"""
    async with semaphore:
        messages_in_flight.inc()
        start_time = asyncio.get_event_loop().time()
        endpoint = "unknown"
        
        try:
            # Decode message
            payload = json.loads(msg.data.decode())
            endpoint = payload.get("endpoint", payload.get("original_path", "unknown"))
            http_method = payload.get("method", "POST")  # Get HTTP method from payload
            
            print(f"📨 Processing message for endpoint: {endpoint} (method: {http_method})")
            
            # Forward to external endpoint with correct HTTP method
            api_key = os.getenv("EXTERNAL_ENDPOINT_API_KEY", "")
            success = await forward_to_external(endpoint, payload, api_key if api_key else None, http_method)
            
            duration = asyncio.get_event_loop().time() - start_time
            message_duration.labels(endpoint=endpoint).observe(duration)
            
            if success:
                # Acknowledge message (remove from queue)
                await msg.ack()
                messages_processed.labels(status="success", endpoint=endpoint).inc()
                print(f"✅ Message processed and forwarded successfully for {endpoint}")
            else:
                # Negative acknowledge (requeue for retry)
                await msg.nak()
                messages_processed.labels(status="failed", endpoint=endpoint).inc()
                print(f"❌ Failed to forward message for {endpoint}, requeued for retry")
                
        except json.JSONDecodeError as e:
            print(f"❌ Invalid JSON in message: {e}")
            await msg.term()  # Terminate (don't retry)
            messages_processed.labels(status="invalid", endpoint=endpoint).inc()
        except Exception as e:
            print(f"❌ Error processing message for {endpoint}: {e}")
            await msg.nak()  # Requeue for retry
            messages_processed.labels(status="error", endpoint=endpoint).inc()
        finally:
            messages_in_flight.dec()

async def consume_messages():
    """Consume messages from NATS JetStream"""
    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=-1  # Retry forever
        )
        js = nc.jetstream()
        print("✅ Connected to NATS JetStream")
        
        # Subscribe to messages (using wildcard to catch all API subjects)
        print(f"Subscribing to stream '{NATS_STREAM}', consumer '{NATS_CONSUMER}'...")
        print(f"   Subject pattern: {NATS_SUBJECT}")
        print(f"   Endpoints configured: {len(ENDPOINT_MAPPING)}")
        
        psub = await js.pull_subscribe(
            NATS_SUBJECT,
            NATS_CONSUMER,
            stream=NATS_STREAM
        )
        print(f"✅ Subscribed, waiting for messages...")
        
        # Start metrics server
        start_http_server(9090)
        print("✅ Metrics server started on port 9090")
        
        # Process messages in a loop
        while running:
            try:
                # Fetch messages (batch of up to 10)
                msgs = await psub.fetch(10, timeout=5)
                for msg in msgs:
                    # Process asynchronously
                    asyncio.create_task(process_message(msg))
            except asyncio.TimeoutError:
                # No messages available, continue
                continue
            except Exception as e:
                print(f"⚠️  Error fetching messages: {e}")
                await asyncio.sleep(1)
        
    except Exception as e:
        print(f"❌ Fatal error: {e}")
        raise
    finally:
        if nc:
            await nc.close()
            print("NATS connection closed")

async def main():
    """Main entry point"""
    print("🚀 Starting NATS Worker (Multi-Endpoint)...")
    print(f"   External Base URL: {BASE_URL}")
    print(f"   Concurrency: {WORKER_CONCURRENCY}")
    print(f"   Retry Attempts: {RETRY_ATTEMPTS}")
    print(f"   Configured Endpoints:")
    for fastapi_endpoint, external_url in ENDPOINT_MAPPING.items():
        print(f"     {fastapi_endpoint} → {external_url}")
    
    try:
        await consume_messages()
    except KeyboardInterrupt:
        print("\n🛑 Worker stopped by user")
    except Exception as e:
        print(f"❌ Worker crashed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())