import httpx
import re
from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
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"
}

_http_client: Optional[httpx.AsyncClient] = None

def get_http_client() -> httpx.AsyncClient:
    global _http_client
    if _http_client is None or _http_client.is_closed:
        _http_client = httpx.AsyncClient(
            timeout=httpx.Timeout(connect=5.0, read=120.0, write=60.0, pool=30.0),
            follow_redirects=False,
            limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)
        )
    return _http_client

async def close_http_client():
    global _http_client
    if _http_client and not _http_client.is_closed:
        await _http_client.aclose()
        _http_client = 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}/
    to strictly isolate session cookies between Hermes universes.
    """
    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]:
    client = get_http_client()
    try:
        res = await client.get(backend_url, timeout=3.0)
        return {
            "status": "online" if res.status_code < 500 else "degraded",
            "status_code": res.status_code,
            "latency_ms": int(res.elapsed.total_seconds() * 1000)
        }
    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:
    client = get_http_client()
    
    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 = {}
    for key, value in request.headers.items():
        if key.lower() not in HOP_BY_HOP_HEADERS and key.lower() != "host":
            req_headers[key] = value
    
    body = await request.body()
    
    try:
        upstream_req = client.build_request(
            method=request.method,
            url=target_url,
            headers=req_headers,
            content=body
        )
        
        upstream_res = await client.send(upstream_req, stream=True)
        
        # 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 universe_id:
                # Cloisonnement strict des cookies de session par univers
                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 universe_id:
                # Réécriture de la redirection si le backend renvoie vers la racine /
                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"
        )