"""
Interceptor Module - Transparent request proxy for custom user services

Provides a transparent proxy (like nginx) that forwards requests from
IEOPS proxy to user's custom service implementation.

Note: For unified management of Register + Interceptor, use Engine instead.

Usage:
```python
from gmi_ieops.handler import Engine
from gmi_ieops.utils import get_service_port

port = get_service_port()
engine = Engine(service_port=port)
engine.start()
app.run(host="127.0.0.1", port=port)
```
"""

import os
import asyncio
import time
from typing import Optional, Set

import aiohttp
from aiohttp import web

from ..utils.log import log
from ..utils.util import get_socket_path
from ..config import env
from .load import get_load_calculator
from .billing import get_billing_tracker, UsageExtractor

# Default headers to filter when proxying requests (client -> upstream)
DEFAULT_FILTER_HEADERS = frozenset({
    "host",
    "connection",
    "keep-alive",
    "transfer-encoding",
    "origin",
    "referer",
})

# Default headers to filter when proxying responses (upstream -> client)
DEFAULT_FILTER_RESPONSE_HEADERS = frozenset({
    "transfer-encoding",
    "connection",
    "content-encoding",
    "content-length",
    "content-type",  # handled separately via content_type/charset params
})


class Interceptor:
    """
    Transparent Request Proxy using aiohttp
    
    Pure proxy - no business logic, no Register management.
    For Register + Interceptor, use Engine instead.
    
    Features:
    - HTTP request forwarding (all methods)
    - WebSocket bidirectional proxy
    - SSE/streaming response support
    """
    
    def __init__(
        self,
        upstream: str,
        timeout: int = 600,
        app_name: Optional[str] = None,
        socket_path: Optional[str] = None,
        filter_headers: Optional[Set[str]] = None,
        filter_response_headers: Optional[Set[str]] = None,
    ) -> None:
        """
        Initialize interceptor.
        
        Args:
            upstream: Upstream service URL (e.g., "http://127.0.0.1:8080")
            timeout: Request timeout in seconds
            app_name: Application name (for logging)
            socket_path: Unix socket path for listening
            filter_headers: Set of header names (lowercase) to filter from requests.
                           Defaults to DEFAULT_FILTER_HEADERS.
            filter_response_headers: Set of header names (lowercase) to filter from responses.
                           Defaults to DEFAULT_FILTER_RESPONSE_HEADERS.
        """
        self._upstream = upstream.rstrip("/")
        self._timeout = aiohttp.ClientTimeout(total=timeout)
        self._app_name = app_name or env.app.name
        self._socket_path = socket_path or get_socket_path()
        self._filter_headers = filter_headers if filter_headers is not None else DEFAULT_FILTER_HEADERS
        self._filter_response_headers = filter_response_headers if filter_response_headers is not None else DEFAULT_FILTER_RESPONSE_HEADERS
        
        self._session: Optional[aiohttp.ClientSession] = None
    
    @property
    def socket_path(self) -> str:
        """Get socket path"""
        return self._socket_path
    
    async def _on_startup(self, app: web.Application):
        """Startup handler"""
        log.get_logger().info(f"{self._app_name} Interceptor starting...")
        log.get_logger().info(f"Upstream: {self._upstream}")
        log.get_logger().info(f"Socket: {self._socket_path}")
        
        # 配置连接池：缩短 keepalive 避免 stale connection，启用连接复用
        connector = aiohttp.TCPConnector(
            limit=100,
            limit_per_host=0,  # 单 host 不限
            keepalive_timeout=10,  # 10s keepalive（默认 15s），减少 stale 概率
            enable_cleanup_closed=True,  # 主动清理已关闭连接
        )
        self._session = aiohttp.ClientSession(timeout=self._timeout, connector=connector)
    
    async def _on_shutdown(self, app: web.Application):
        """Shutdown handler"""
        log.get_logger().info(f"{self._app_name} Interceptor shutting down...")
        
        if self._session:
            await self._session.close()
        
        if self._socket_path and os.path.exists(self._socket_path):
            try:
                os.remove(self._socket_path)
            except Exception:
                pass
    
    async def _proxy_http(self, request: web.Request, org_id: str = "",
                          model: str = "") -> web.StreamResponse:
        """Proxy HTTP request, optionally extracting token usage for billing."""
        # Use path_qs to preserve URL encoding (e.g., gmifs%3A%2F%2F stays encoded)
        # request.path would decode it to gmifs:// which breaks routing
        url = f"{self._upstream}{request.path_qs}"
        trace_id = request.headers.get("X-Trace-Id", "-")
        
        # Filter out headers that shouldn't be forwarded
        headers = {
            k: v for k, v in request.headers.items()
            if k.lower() not in self._filter_headers
        }
        
        body = await request.read()
        
        # 重试一次：aiohttp 对 POST 请求不自动重试 stale connection，
        # 手动重试可以覆盖 connection reset / server disconnect 场景
        last_err = None
        for attempt in range(2):
            try:
                return await self._do_proxy_http(request, url, headers, body, trace_id,
                                                 org_id=org_id, model=model)
            except aiohttp.ClientError as e:
                last_err = e
                if attempt == 0:
                    log.get_logger(trace_id).warning(f"Upstream connection error (will retry): {e} path={request.path}")
                    continue
            except asyncio.TimeoutError:
                log.get_logger(trace_id).error(f"502: upstream timeout path={request.path}")
                return web.Response(text="Gateway Timeout", status=504)
        
        log.get_logger(trace_id).error(f"502: upstream error after retry: {last_err} path={request.path}")
        return web.Response(text=f"Bad Gateway: {last_err}", status=502)
    
    async def _do_proxy_http(self, request: web.Request, url: str, headers: dict,
                             body: bytes, trace_id: str, org_id: str = "",
                             model: str = "") -> web.StreamResponse:
        """Execute single HTTP proxy attempt. Raises aiohttp.ClientError on connection failure."""
        billing_enabled = env.billing.enabled and (org_id or model)

        async with self._session.request(
            method=request.method,
            url=url,
            headers=headers,
            data=body if body else None,
        ) as resp:
            # 直接用 aiohttp 已解析好的 content_type/charset，避免手动 split 字符串
            parsed_content_type = resp.content_type or "application/octet-stream"
            charset = resp.charset
            is_stream = (
                parsed_content_type == "text/event-stream"
                or "chunked" in resp.headers.get("Transfer-Encoding", "")
            )
            resp_headers = {
                k: v for k, v in resp.headers.items()
                if k.lower() not in self._filter_response_headers
            }

            if is_stream:
                response = web.StreamResponse(status=resp.status, headers=resp_headers)
                response.content_type = parsed_content_type
                if charset:
                    response.charset = charset
                await response.prepare(request)

                # SSE 按行迭代确保事件即时转发；其他流用 iter_any() 提高性能
                if parsed_content_type == "text/event-stream":
                    # Only extract billing from successful SSE streams
                    track_billing = billing_enabled and resp.status == 200
                    usage_record = None
                    async for line in resp.content:
                        await response.write(line)
                        # Extract usage from SSE (typically in the last chunk)
                        if track_billing:
                            try:
                                decoded = line.decode("utf-8", errors="replace").strip()
                                extracted = UsageExtractor.extract_from_sse_line(decoded)
                                if extracted:
                                    usage_record = extracted
                            except Exception:
                                pass  # billing extraction must never break the proxy
                    # Record accumulated SSE usage
                    if usage_record:
                        try:
                            get_billing_tracker().record(org_id=org_id, model=model, usage=usage_record)
                        except Exception:
                            pass
                else:
                    async for chunk in resp.content.iter_any():
                        await response.write(chunk)

                await response.write_eof()
                return response
            else:
                content = await resp.read()
                # Extract usage from non-streaming JSON response
                if billing_enabled and resp.status == 200:
                    try:
                        usage_record = UsageExtractor.extract_from_json(content)
                        if usage_record:
                            get_billing_tracker().record(org_id=org_id, model=model, usage=usage_record)
                    except Exception:
                        pass  # billing extraction must never break the proxy
                return web.Response(
                    body=content,
                    status=resp.status,
                    headers=resp_headers,
                    content_type=parsed_content_type,
                    charset=charset,
                )
    
    async def _proxy_websocket(self, request: web.Request) -> web.WebSocketResponse:
        """Proxy WebSocket"""
        client_ws = web.WebSocketResponse()
        await client_ws.prepare(request)
        
        ws_url = self._upstream.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}{request.path}"
        
        try:
            async with self._session.ws_connect(ws_url) as upstream_ws:
                async def c2u():
                    async for msg in client_ws:
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            await upstream_ws.send_str(msg.data)
                        elif msg.type == aiohttp.WSMsgType.BINARY:
                            await upstream_ws.send_bytes(msg.data)
                        elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                            break
                
                async def u2c():
                    async for msg in upstream_ws:
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            await client_ws.send_str(msg.data)
                        elif msg.type == aiohttp.WSMsgType.BINARY:
                            await client_ws.send_bytes(msg.data)
                        elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                            break
                
                await asyncio.gather(c2u(), u2c(), return_exceptions=True)
        except Exception as e:
            log.get_logger().error(f"WebSocket proxy error: {e}")
        
        return client_ws
    
    async def _handle_request(self, request: web.Request) -> web.StreamResponse:
        """Route handler - tracks load for each proxied request"""
        trace_id = request.headers.get("X-Trace-Id", "-")
        start = time.monotonic()
        calc = get_load_calculator()
        calc.on_request_start()
        try:
            if (
                request.headers.get("Upgrade", "").lower() == "websocket" and
                request.headers.get("Connection", "").lower() == "upgrade"
            ):
                return await self._proxy_websocket(request)
            # Extract billing context from request headers
            org_id = request.headers.get(env.billing.org_id_header, "")
            model = request.headers.get(env.billing.model_header, "")
            resp = await self._proxy_http(request, org_id=org_id, model=model)
            latency = time.monotonic() - start
            log.get_logger(trace_id).info(
                f"{request.method} {request.path} status={resp.status} latency={latency:.3f}s"
            )
            return resp
        finally:
            calc.on_request_end()
    
    def serve(self) -> None:
        """Start server (blocking)"""
        app = web.Application()
        app.on_startup.append(self._on_startup)
        app.on_shutdown.append(self._on_shutdown)
        app.router.add_route("*", "/{path:.*}", self._handle_request)
        app.router.add_route("*", "/", self._handle_request)
        
        web.run_app(
            app,
            path=self._socket_path,
            print=lambda *args: None,
        )
    
    async def serve_async(self) -> web.AppRunner:
        """Start server (async, non-blocking)"""
        app = web.Application()
        app.on_startup.append(self._on_startup)
        app.on_shutdown.append(self._on_shutdown)
        app.router.add_route("*", "/{path:.*}", self._handle_request)
        app.router.add_route("*", "/", self._handle_request)
        
        runner = web.AppRunner(app)
        await runner.setup()
        
        site = web.UnixSite(runner, self._socket_path)
        await site.start()
        
        return runner
