import os
from ..utils.util import randstr, APP_ID
from ..utils.log import *
import aiohttp
import asyncio
import requests
from typing import Dict, Any, TypedDict, Optional, Literal
from urllib.parse import quote
import time

_APP_NAME = os.getenv("APP_NAME", "ieops")
_SERVER_SOCKET_DIR = os.getenv("SERVER_SOCKET_DIR", os.getcwd()) # 默认是项目根目录

_MODEL_THREAD_CONCURRENCY = os.getenv("MODEL_THREAD_CONCURRENCY", "8")
_REGISTER_INTERVAL = os.getenv("REGISTER_INTERVAL", "10")
_REGISTER_ENABLED = os.getenv("REGISTER_ENABLED", "true")

# connector type
ConnectorType = Literal["proxy", "brslet"]

# socket 配置
_SOCKET_PATHS: Dict[ConnectorType, str] = {
    "proxy": os.getenv("IEOPS_PROXY_SOCKET", "/var/run/ieops/proxy.sock"),
    "brslet": os.getenv("IEOPS_BRSLET_SOCKET", "/var/run/ieops/brslet.sock"),
}

# async unix connectors (lazy init)
_CONNECTORS: Dict[ConnectorType, Optional[aiohttp.UnixConnector]] = {
    "proxy": None,
    "brslet": None,
}

# sync session (lazy init)
_SYNC_SESSION: Optional[requests.Session] = None


def _build_unix_url(socket_path: str, api_path: str) -> str:
    """构建 requests_unixsocket 格式的 Unix socket URL
    
    Args:
        socket_path: socket 文件路径，如 /var/run/ieops/proxy.sock
        api_path: API 路径，如 register
        
    Returns:
        http+unix://%2Fvar%2Frun%2Fieops%2Fproxy.sock/register
    """
    encoded_socket = quote(socket_path, safe='')
    return f"http+unix://{encoded_socket}/{api_path}"


class Payload(TypedDict):
    path: str
    payload: Dict[str, Any]

class Response:
    """HTTP 响应封装"""
    def __init__(self, message: str, data: Optional[Any] = None):
        self.message = message
        self.data = data
    
    def __repr__(self) -> str:
        return f"Response(message={self.message!r}, data={self.data!r})"


def _get_connector(connector_type: ConnectorType) -> aiohttp.UnixConnector:
    """获取指定类型的异步 Unix socket connector（懒加载）"""
    if _CONNECTORS[connector_type] is None:
        _CONNECTORS[connector_type] = aiohttp.UnixConnector(path=_SOCKET_PATHS[connector_type])
    return _CONNECTORS[connector_type]


def _get_sync_session() -> requests.Session:
    """获取同步 HTTP session，支持 Unix socket"""
    global _SYNC_SESSION
    if _SYNC_SESSION is None:
        _SYNC_SESSION = requests.Session()
        # 添加 Unix socket 支持
        try:
            from requests_unixsocket import UnixAdapter
            _SYNC_SESSION.mount("http+unix://", UnixAdapter())
        except ImportError:
            uvicorn_logger.warning("requests_unixsocket not installed, Unix socket support disabled")
    return _SYNC_SESSION


def post(_payload: Payload, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """send POST request via Unix socket
    
    Args:
        _payload: 请求数据
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    session = _get_sync_session()
    url = _build_unix_url(_SOCKET_PATHS[connector_type], _payload['path'])
    for attempt in range(max_retries + 1):
        try:
            resp = session.post(url, json=_payload['payload'], timeout=timeout)
            if resp.status_code == 200:
                try:
                    data = resp.json()
                    return Response(message=data.get('message', ''), data=data.get('data'))
                except ValueError as e:
                    err = Response(message=f"Invalid JSON response: {e}")
            else:
                err = Response(message=f"HTTP {resp.status_code}: {resp.text[:200]}")
        except requests.Timeout as e:
            err = Response(message=f"Client POST [{_payload['path']}] timeout: {e}")
        except requests.RequestException as e:
            err = Response(message=f"Client POST [{_payload['path']}] failed: {e}")
        if attempt < max_retries:
            time.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client POST [{_payload['path']}] failed after {max_retries + 1} attempts: {err}")
    return err


def get(path: str, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """send GET request via Unix socket
    
    Args:
        path: API 路径
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    session = _get_sync_session()
    url = _build_unix_url(_SOCKET_PATHS[connector_type], path)
    for attempt in range(max_retries + 1):
        try:
            resp = session.get(url, timeout=timeout)
            if resp.status_code == 200:
                try:
                    data = resp.json()
                    return Response(message=data.get('message', ''), data=data.get('data'))
                except ValueError as e:
                    err = Response(message=f"Invalid JSON response: {e}")
            else:
                err = Response(message=f"HTTP {resp.status_code}: {resp.text[:200]}")
        except requests.Timeout as e:
            err = Response(message=f"Client GET [{path}] timeout: {e}")
        except requests.RequestException as e:
            err = Response(message=f"Client GET [{path}] failed: {e}")
        if attempt < max_retries:
            time.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client GET [{path}] failed after {max_retries + 1} attempts: {err}")
    return err


def delete(path: str, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """send DELETE request via Unix socket
    
    Args:
        path: API 路径
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    session = _get_sync_session()
    url = _build_unix_url(_SOCKET_PATHS[connector_type], path)
    for attempt in range(max_retries + 1):
        try:
            resp = session.delete(url, timeout=timeout)
            if resp.status_code == 200:
                try:
                    data = resp.json()
                    return Response(message=data.get('message', ''), data=data.get('data'))
                except ValueError as e:
                    err = Response(message=f"Invalid JSON response: {e}")
            else:
                err = Response(message=f"HTTP {resp.status_code}: {resp.text[:200]}")
        except requests.Timeout as e:
            err = Response(message=f"Client DELETE [{path}] timeout: {e}")
        except requests.RequestException as e:
            err = Response(message=f"Client DELETE [{path}] failed: {e}")
        if attempt < max_retries:
            time.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client DELETE [{path}] failed after {max_retries + 1} attempts: {err}")
    return err


async def async_post(_payload: Payload, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """async POST request via Unix socket
    
    Args:
        _payload: 请求数据
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    url = f"http://localhost/{_payload['path']}"  # aiohttp UnixConnector 使用 localhost
    client_timeout = aiohttp.ClientTimeout(total=timeout)
    
    for attempt in range(max_retries + 1):
        try:
            async with aiohttp.ClientSession(
                connector=_get_connector(connector_type), 
                connector_owner=False,
                timeout=client_timeout
            ) as session:
                async with session.post(url, json=_payload['payload']) as resp:
                    if resp.status == 200:
                        try:
                            data = await resp.json()
                            return Response(message=data.get('message', ''), data=data.get('data'))
                        except (ValueError, aiohttp.ContentTypeError) as e:
                            err = Response(message=f"Invalid JSON response: {e}")
                    else:
                        text = await resp.text()
                        err = Response(message=f"HTTP {resp.status}: {text[:200]}")
        except asyncio.TimeoutError:
            err = Response(message=f"Client POST [{_payload['path']}] timeout")
        except aiohttp.ClientError as e:
            err = Response(message=f"Client POST [{_payload['path']}] failed: {e}")
        except asyncio.CancelledError:
            raise  # 重新抛出，让调用者处理取消
        if attempt < max_retries:
            await asyncio.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client POST [{_payload['path']}] failed after {max_retries + 1} attempts: {err}")
    return err


async def async_get(path: str, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """async GET request via Unix socket
    
    Args:
        path: API 路径
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    url = f"http://localhost/{path}"  # aiohttp UnixConnector 使用 localhost
    client_timeout = aiohttp.ClientTimeout(total=timeout)
    
    for attempt in range(max_retries + 1):
        try:
            async with aiohttp.ClientSession(
                connector=_get_connector(connector_type), 
                connector_owner=False,
                timeout=client_timeout
            ) as session:
                async with session.get(url) as resp:
                    if resp.status == 200:
                        try:
                            data = await resp.json()
                            return Response(message=data.get('message', ''), data=data.get('data'))
                        except (ValueError, aiohttp.ContentTypeError) as e:
                            err = Response(message=f"Invalid JSON response: {e}")
                    else:
                        text = await resp.text()
                        err = Response(message=f"HTTP {resp.status}: {text[:200]}")
        except asyncio.TimeoutError:
            err = Response(message=f"Client GET [{path}] timeout")
        except aiohttp.ClientError as e:
            err = Response(message=f"Client GET [{path}] failed: {e}")
        except asyncio.CancelledError:
            raise  # 重新抛出，让调用者处理取消
        if attempt < max_retries:
            await asyncio.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client GET [{path}] failed after {max_retries + 1} attempts: {err}")
    return err


async def async_delete(path: str, timeout: float = 0.8, max_retries: int = 2, connector_type: ConnectorType = "proxy") -> Response:
    """async DELETE request via Unix socket
    
    Args:
        path: API 路径
        timeout: 超时时间
        max_retries: 最大重试次数
        connector_type: "proxy" 或 "brslet"
    """
    err = None
    url = f"http://localhost/{path}"  # aiohttp UnixConnector 使用 localhost
    client_timeout = aiohttp.ClientTimeout(total=timeout)
    
    for attempt in range(max_retries + 1):
        try:
            async with aiohttp.ClientSession(
                connector=_get_connector(connector_type), 
                connector_owner=False,
                timeout=client_timeout
            ) as session:
                async with session.delete(url) as resp:
                    if resp.status == 200:
                        try:
                            data = await resp.json()
                            return Response(message=data.get('message', ''), data=data.get('data'))
                        except (ValueError, aiohttp.ContentTypeError) as e:
                            err = Response(message=f"Invalid JSON response: {e}")
                    else:
                        text = await resp.text()
                        err = Response(message=f"HTTP {resp.status}: {text[:200]}")
        except asyncio.TimeoutError:
            err = Response(message=f"Client DELETE [{path}] timeout")
        except aiohttp.ClientError as e:
            err = Response(message=f"Client DELETE [{path}] failed: {e}")
        except asyncio.CancelledError:
            raise  # 重新抛出，让调用者处理取消
        if attempt < max_retries:
            await asyncio.sleep(0.1 * (attempt + 1))
        else:
            uvicorn_logger.warning(f"Client DELETE [{path}] failed after {max_retries + 1} attempts: {err}")
    return err

class Register:
    def __init__(self):
        self._appinfo = self._app_info()
        
    def _app_info(self):
        self._tokens = [randstr(8) for _ in range(int(_MODEL_THREAD_CONCURRENCY))]
        return {
            "id": APP_ID,
            "server_socket": f"{_SERVER_SOCKET_DIR}/{APP_ID}.sock",
            "max_concurrent_reqs": int(_MODEL_THREAD_CONCURRENCY),
            "endpoint": _APP_NAME,
            "weight": 1,
        }

    async def _unregister_old(self):
        """启动时先取消旧的注册，清理可能存在的僵尸注册
        按 socket 目录清理，避免误伤同 endpoint 但不同 pod 的服务
        """
        socket_dir = _SERVER_SOCKET_DIR
        uvicorn_logger.info("Unregistering old models in socket_dir {} from IEOPS proxy (if exists)...".format(socket_dir))
        # 取消同一目录下的旧注册（同一 pod 重启后目录相同，但 socket ID 不同）
        await async_post(Payload(path="unregister", payload={"socket_dir": socket_dir}), max_retries=0)

    async def _unregister(self):
        """停止时取消当前注册"""
        uvicorn_logger.info("Unregistering {} from IEOPS proxy...".format(_APP_NAME))
        await async_post(Payload(path="unregister", payload={"id": APP_ID}), max_retries=1)

    async def register(self):
        if _REGISTER_ENABLED.lower() == "false":
            uvicorn_logger.info("Register is disabled, skipping...")
            return
        uvicorn_logger.info("Registering {} to IEOPS proxy...".format(_APP_NAME))
        try:
            # 方案3: 启动时先清理旧注册
            await self._unregister_old()
            while True:
                await async_post(Payload(path="register", payload=self._appinfo))
                await asyncio.sleep(int(_REGISTER_INTERVAL))
        except asyncio.CancelledError:
            uvicorn_logger.info("Register stopped")
            raise
        finally:
            # 停止时取消注册
            try:
                await self._unregister()
            except Exception as e:
                uvicorn_logger.warning(f"Failed to unregister: {e}")
            # 关闭所有 connectors
            for connector in _CONNECTORS.values():
                if connector is not None:
                    await connector.close()
            if _SYNC_SESSION is not None:
                _SYNC_SESSION.close()

__all__ = [
    'Payload', 'Response', 'ConnectorType',
    'post', 'get', 'delete', 
    'async_post', 'async_get', 'async_delete', 
    'Register',
]