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",
    "origin"
}

_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:
    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}"
    
    is_dsh = "dsh-vps:3080" in backend_url or universe_id == "dsh"
    
    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 in HOP_BY_HOP_HEADERS or k_str == "host":
            continue
        if is_dsh and k_str == "origin":
            continue
        req_headers.append((raw_k, raw_v))
            
    # If DSH backend, replicate localhost:3080 Host/Origin for internal origin check
    if is_dsh:
        req_headers.append((b"host", b"localhost:3080"))
        req_headers.append((b"origin", b"http://localhost:3080"))
    else:
        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
        )
        
        upstream_res = await transport.handle_async_request(upstream_req)
        
        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:
                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:
                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 réseau interne 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
):
    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}"
        
    is_dsh = "dsh-vps:3080" in backend_url or universe_id == "dsh"
    
    upstream_headers = {}
    for k, v in client_ws.headers.items():
        if k.lower() not in WS_HOP_BY_HOP_HEADERS:
            upstream_headers[k] = v
            
    if is_dsh:
        upstream_headers["host"] = "localhost:3080"
        origin_val = "http://localhost:3080"
    else:
        upstream_headers["x-forwarded-host"] = client_ws.headers.get("host", "")
        upstream_headers["x-forwarded-proto"] = client_ws.url.scheme or "http"
        origin_val = None
        
    if client_ws.client:
        upstream_headers["x-forwarded-for"] = client_ws.client.host
        
    subprotocols_raw = client_ws.headers.get("sec-websocket-protocol", "")
    subprotocols = [s.strip() for s in subprotocols_raw.split(",") if s.strip()] or None
    
    try:
        connect_kwargs = {
            "additional_headers": upstream_headers,
            "subprotocols": subprotocols,
            "ping_interval": 20,
            "ping_timeout": 20,
            "max_size": 10 * 1024 * 1024
        }
        if origin_val:
            connect_kwargs["origin"] = origin_val
            
        async with websockets.connect(upstream_url, **connect_kwargs) as upstream_ws:
            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