import httpx
import re
import asyncio
from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
from fastapi import Request, Response, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse
import websockets
from websockets.exceptions import ConnectionClosed
import logging

logger = logging.getLogger("hermes_hub.proxy")

HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "content-length"
}

WS_HOP_BY_HOP_HEADERS = {
    "connection",
    "upgrade",
    "sec-websocket-key",
    "sec-websocket-version",
    "sec-websocket-extensions",
    "sec-websocket-accept",
    "host"
}

_http_transport: Optional[httpx.AsyncHTTPTransport] = None

def get_http_transport() -> httpx.AsyncHTTPTransport:
    global _http_transport
    if _http_transport is None:
        _http_transport = httpx.AsyncHTTPTransport(
            retries=1,
            limits=httpx.Limits(max_keepalive_connections=50, max_connections=100, keepalive_expiry=30.0)
        )
    return _http_transport

async def close_http_client():
    global _http_transport
    if _http_transport:
        await _http_transport.aclose()
        _http_transport = None

def rewrite_cookie_path(cookie_header: str, universe_id: str) -> str:
    """
    Rewrites the Path attribute of a Set-Cookie header to /u/{universe_id}/
    when using path-based proxying.
    """
    target_path = f"/u/{universe_id}/"
    if re.search(r'(?i)\bpath=[^;]*', cookie_header):
        return re.sub(r'(?i)\bpath=[^;]*', f'Path={target_path}', cookie_header)
    else:
        return f"{cookie_header}; Path={target_path}"

async def check_backend_health(backend_url: str) -> Dict[str, Any]:
    transport = get_http_transport()
    try:
        req = httpx.Request("GET", backend_url)
        res = await transport.handle_async_request(req)
        status_code = res.status_code
        await res.aclose()
        return {
            "status": "online" if status_code < 500 else "degraded",
            "status_code": status_code,
            "latency_ms": 100
        }
    except Exception as e:
        return {
            "status": "offline",
            "error": str(e),
            "latency_ms": None
        }

async def proxy_request(
    request: Request,
    backend_url: str,
    path: str,
    universe_id: Optional[str] = None
) -> Response:
    transport = get_http_transport()
    
    base_url = backend_url.rstrip("/")
    sub_path = path.lstrip("/")
    target_url = f"{base_url}/{sub_path}" if sub_path else base_url
    
    if request.url.query:
        target_url = f"{target_url}?{request.url.query}"
    
    req_headers: List[Tuple[bytes, bytes]] = []
    for raw_k, raw_v in request.headers.raw:
        k_str = raw_k.decode("latin-1").lower()
        if k_str not in HOP_BY_HOP_HEADERS and k_str != "host":
            req_headers.append((raw_k, raw_v))
            
    # Forward original host header info
    req_headers.append((b"x-forwarded-host", request.headers.get("host", "").encode("latin-1")))
    req_headers.append((b"x-forwarded-proto", (request.url.scheme or "http").encode("latin-1")))
    
    body = await request.body()
    
    is_path_proxied = universe_id and request.url.path.startswith(f"/u/{universe_id}")
    
    try:
        upstream_req = httpx.Request(
            method=request.method,
            url=target_url,
            headers=req_headers,
            content=body
        )
        
        # Pure stateless async request execution — NO cookie jar retention
        upstream_res = await transport.handle_async_request(upstream_req)
        
        # Build raw headers list for precise multi-header control
        raw_headers: List[Tuple[bytes, bytes]] = []
        media_type = upstream_res.headers.get("content-type")
        
        for raw_k, raw_v in upstream_res.headers.raw:
            k_str = raw_k.decode("latin-1").lower()
            v_str = raw_v.decode("latin-1")
            
            if k_str in HOP_BY_HOP_HEADERS:
                continue
                
            if k_str == "set-cookie" and is_path_proxied:
                # Path-based proxying: rewrite cookie path
                v_str = rewrite_cookie_path(v_str, universe_id)
                raw_headers.append((b"set-cookie", v_str.encode("latin-1")))
            elif k_str == "location" and is_path_proxied:
                # Path-based proxying: rewrite location redirect
                if v_str.startswith("/") and not v_str.startswith(f"/u/{universe_id}"):
                    v_str = f"/u/{universe_id}{v_str}"
                raw_headers.append((b"location", v_str.encode("latin-1")))
            else:
                raw_headers.append((raw_k, v_str.encode("latin-1")))
        
        async def stream_content() -> AsyncGenerator[bytes, None]:
            try:
                async for chunk in upstream_res.aiter_raw():
                    yield chunk
            finally:
                await upstream_res.aclose()
        
        response = StreamingResponse(
            stream_content(),
            status_code=upstream_res.status_code,
            media_type=media_type
        )
        response.raw_headers = raw_headers
        return response

    except httpx.ConnectError:
        logger.error(f"Failed to connect to backend at {target_url}")
        return Response(
            content=f"<html><head><title>Backend Unavailable</title></head><body style='background:#121212;color:#ef4444;font-family:sans-serif;padding:2rem;'><h2>Instance Hermes inaccessible</h2><p>Le backend sur <code>{backend_url}</code> ne répond pas. Vérifiez le tunnel Tailscale ou l'état du conteneur.</p></body></html>",
            status_code=502,
            media_type="text/html"
        )
    except Exception as e:
        logger.exception(f"Proxy error for {target_url}: {e}")
        return Response(
            content=f"Proxy Error: {str(e)}",
            status_code=500,
            media_type="text/plain"
        )

async def proxy_websocket(
    client_ws: WebSocket,
    backend_url: str,
    path: str,
    universe_id: Optional[str] = None
):
    """
    Bi-directional full duplex WebSocket proxy bridge.
    Forwards connection from client browser to backend Hermes universe instance.
    """
    ws_base = backend_url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/")
    sub_path = path.lstrip("/")
    upstream_url = f"{ws_base}/{sub_path}" if sub_path else ws_base
    
    if client_ws.url.query:
        upstream_url = f"{upstream_url}?{client_ws.url.query}"
        
    upstream_headers = {}
    for k, v in client_ws.headers.items():
        if k.lower() not in WS_HOP_BY_HOP_HEADERS:
            upstream_headers[k] = v
            
    upstream_headers["x-forwarded-host"] = client_ws.headers.get("host", "")
    upstream_headers["x-forwarded-proto"] = client_ws.url.scheme or "http"
    if client_ws.client:
        upstream_headers["x-forwarded-for"] = client_ws.client.host
        
    # Check subprotocols
    subprotocols_raw = client_ws.headers.get("sec-websocket-protocol", "")
    subprotocols = [s.strip() for s in subprotocols_raw.split(",") if s.strip()] or None
    
    try:
        async with websockets.connect(
            upstream_url,
            extra_headers=upstream_headers,
            subprotocols=subprotocols,
            ping_interval=20,
            ping_timeout=20,
            max_size=10 * 1024 * 1024
        ) as upstream_ws:
            # Accept client connection with negotiated subprotocol
            await client_ws.accept(subprotocol=upstream_ws.subprotocol)
            
            async def client_to_upstream():
                try:
                    while True:
                        msg = await client_ws.receive()
                        if "text" in msg:
                            await upstream_ws.send(msg["text"])
                        elif "bytes" in msg:
                            await upstream_ws.send(msg["bytes"])
                        elif msg.get("type") == "websocket.disconnect":
                            break
                except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
                    pass
                    
            async def upstream_to_client():
                try:
                    async for msg in upstream_ws:
                        if isinstance(msg, str):
                            await client_ws.send_text(msg)
                        elif isinstance(msg, bytes):
                            await client_ws.send_bytes(msg)
                except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
                    pass
                    
            done, pending = await asyncio.wait(
                [
                    asyncio.create_task(client_to_upstream()),
                    asyncio.create_task(upstream_to_client())
                ],
                return_when=asyncio.FIRST_COMPLETED
            )
            for task in pending:
                task.cancel()
                
    except (WebSocketDisconnect, ConnectionClosed):
        pass
    except Exception as e:
        logger.error(f"WebSocket proxy bridge error for {upstream_url}: {e}")
        try:
            await client_ws.close(code=1011)
        except Exception:
            pass