"""
ComfyUI HTTPS Bridge
====================
在 ComfyUI 启动时自动开启一个本地 HTTPS 代理服务器，
让部署在公网 HTTPS 页面的套壳应用可以安全连接到本地 ComfyUI。

安装：将本文件夹放入 ComfyUI/custom_nodes/ 目录，重启 ComfyUI 即可。

地址检测优先级：
  1. 环境变量 COMFY_HOST / COMFY_PORT（手动覆盖，优先级最高）
  2. ComfyUI 启动参数 --listen / --port（自动读取）
  3. 默认值 127.0.0.1:8188

环境变量（可选）：
  COMFY_BRIDGE_PORT  桥接端口，默认 9999
  COMFY_HOST         覆盖 ComfyUI 地址
  COMFY_PORT         覆盖 ComfyUI 端口
"""

import asyncio
import ssl
import os
import socket
import threading
import ipaddress
import datetime
import pathlib
import traceback

# ── 配置 ──────────────────────────────────────────────────────────────────────

BRIDGE_PORT = int(os.environ.get('COMFY_BRIDGE_PORT', 9999))


def _get_local_ip():
    """获取本机对外的实际 IP 地址（非 127.0.0.1）。"""
    try:
        # 连接一个外部地址（不会真正发包），借此确定本机出口 IP
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return '127.0.0.1'


def _detect_comfy_address():
    """
    按优先级检测 ComfyUI 实际监听地址：
      1. 环境变量 COMFY_HOST / COMFY_PORT
      2. ComfyUI 启动参数 --listen / --port
      3. 默认 127.0.0.1:8188
    """
    host, port = '127.0.0.1', 8188

    # 从 ComfyUI 内部读取启动参数
    try:
        from comfy.cli_args import args
        if args.port:
            port = int(args.port)
        if args.listen:
            listen = args.listen.strip()
            # 0.0.0.0 / :: 表示监听所有接口，取本机实际出口 IP
            is_wildcard = listen in ('0.0.0.0', '::') or '0.0.0.0' in listen or listen.startswith('::')
            host = _get_local_ip() if is_wildcard else listen
    except Exception:
        pass  # 不在 ComfyUI 环境中也能降级运行

    # 环境变量最终覆盖（优先级最高）
    host = os.environ.get('COMFY_HOST', host)
    port = int(os.environ.get('COMFY_PORT', port))

    return host, port


COMFY_HOST, COMFY_PORT = _detect_comfy_address()
CERT_DIR    = pathlib.Path(__file__).parent / '.certs'

CORS_HEADERS = {
    'Access-Control-Allow-Origin':  '*',
    'Access-Control-Allow-Methods': '*',
    'Access-Control-Allow-Headers': '*',
}

# ── 证书生成 ──────────────────────────────────────────────────────────────────

def generate_cert():
    """使用 cryptography 库生成自签名证书，无需任何外部工具。"""
    try:
        from cryptography import x509
        from cryptography.x509.oid import NameOID
        from cryptography.hazmat.primitives import hashes, serialization
        from cryptography.hazmat.primitives.asymmetric import rsa
    except ImportError:
        print('[HTTPS Bridge] ❌ 缺少 cryptography 库，请在 ComfyUI 的 Python 环境中运行：')
        print('               pip install cryptography')
        return None, None

    cert_file = CERT_DIR / 'cert.pem'
    key_file  = CERT_DIR / 'key.pem'

    if cert_file.exists() and key_file.exists():
        return str(cert_file), str(key_file)

    CERT_DIR.mkdir(exist_ok=True)

    # 收集本机所有需要覆盖的 IP
    local_ip = _get_local_ip()
    san_ips = {ipaddress.IPv4Address('127.0.0.1')}
    if local_ip != '127.0.0.1':
        try:
            san_ips.add(ipaddress.IPv4Address(local_ip))
        except ValueError:
            pass

    print(f'[HTTPS Bridge] 🔐 首次运行，生成自签名证书（有效期 10 年）...')
    print(f'               覆盖地址: localhost, 127.0.0.1' + (f', {local_ip}' if local_ip != '127.0.0.1' else ''))

    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)

    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, 'localhost'),
    ])
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
        .add_extension(
            x509.SubjectAlternativeName([
                x509.DNSName('localhost'),
                *[x509.IPAddress(ip) for ip in san_ips],
            ]),
            critical=False,
        )
        .sign(key, hashes.SHA256())
    )

    key_file.write_bytes(key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption(),
    ))
    cert_file.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    print(f'[HTTPS Bridge] ✅ 证书已生成: {CERT_DIR}')
    return str(cert_file), str(key_file)

# ── 请求处理 ──────────────────────────────────────────────────────────────────

def _is_websocket(request):
    return (
        request.headers.get('Upgrade', '').lower() == 'websocket'
        and 'upgrade' in request.headers.get('Connection', '').lower()
    )

def _proxy_headers(headers, skip_keys):
    return {k: v for k, v in headers.items() if k.lower() not in skip_keys}


async def handle_websocket(request):
    """双向代理 WebSocket 连接。"""
    import aiohttp
    from aiohttp import web

    ws_server = web.WebSocketResponse()
    await ws_server.prepare(request)

    target = f'ws://{COMFY_HOST}:{COMFY_PORT}{request.path_qs}'
    skip   = {'host', 'upgrade', 'connection',
               'sec-websocket-key', 'sec-websocket-version', 'sec-websocket-extensions'}
    headers = _proxy_headers(request.headers, skip)

    try:
        async with aiohttp.ClientSession() as session:
            async with session.ws_connect(target, headers=headers) as ws_client:

                async def forward(src, dst, is_text_fn, send_str, send_bytes):
                    async for msg in src:
                        if dst.closed:
                            break
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            await send_str(msg.data)
                        elif msg.type == aiohttp.WSMsgType.BINARY:
                            await send_bytes(msg.data)
                        elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                            break

                await asyncio.gather(
                    forward(ws_client, ws_server, None,
                            ws_server.send_str, ws_server.send_bytes),
                    forward(ws_server, ws_client, None,
                            ws_client.send_str, ws_client.send_bytes),
                    return_exceptions=True,
                )
    except Exception as e:
        print(f'[HTTPS Bridge] WS 错误: {e}')

    return ws_server


async def handle_http(request):
    """代理普通 HTTP 请求，支持流式响应（大文件/图片）。"""
    import aiohttp
    from aiohttp import web

    if request.method == 'OPTIONS':
        return web.Response(status=204, headers=CORS_HEADERS)

    target = f'http://{COMFY_HOST}:{COMFY_PORT}{request.path_qs}'
    skip   = {'host', 'content-length', 'transfer-encoding'}
    headers = _proxy_headers(request.headers, skip)

    try:
        async with aiohttp.ClientSession() as session:
            async with session.request(
                request.method, target,
                headers=headers,
                data=await request.read(),
                allow_redirects=False,
            ) as resp:
                resp_headers = {
                    k: v for k, v in resp.headers.items()
                    if k.lower() not in {'transfer-encoding'}
                }
                resp_headers.update(CORS_HEADERS)

                response = web.StreamResponse(status=resp.status, headers=resp_headers)
                await response.prepare(request)
                async for chunk in resp.content.iter_chunked(65536):
                    await response.write(chunk)
                await response.write_eof()
                return response

    except aiohttp.ClientConnectorError:
        return web.Response(
            status=502,
            text=f'无法连接到 ComfyUI ({COMFY_HOST}:{COMFY_PORT})，请确认 ComfyUI 正在运行。',
            headers=CORS_HEADERS,
        )
    except Exception as e:
        return web.Response(status=500, text=f'代理错误: {e}', headers=CORS_HEADERS)


async def handle_request(request):
    if _is_websocket(request):
        return await handle_websocket(request)
    return await handle_http(request)

# ── 服务器启动 ────────────────────────────────────────────────────────────────

async def _run_server(cert_file, key_file):
    from aiohttp import web

    ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ssl_ctx.load_cert_chain(cert_file, key_file)

    app = web.Application(client_max_size=512 * 1024 * 1024)  # 512MB
    app.router.add_route('*', '/',              handle_request)
    app.router.add_route('*', '/{path_info:.*}', handle_request)

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '0.0.0.0', BRIDGE_PORT, ssl_context=ssl_ctx)
    await site.start()

    print(f'\n✅ [HTTPS Bridge] 已启动')
    print(f'   桥接地址: https://localhost:{BRIDGE_PORT}')
    print(f'   代理目标: http://{COMFY_HOST}:{COMFY_PORT}')
    print(f'\n⚠️  首次使用需在浏览器授权证书（只需一次）：')
    print(f'   1. 打开 https://localhost:{BRIDGE_PORT}')
    print(f'   2. 点击「高级」→「继续访问 localhost」')
    print(f'   3. 看到空白页即表示成功\n')

    await asyncio.Event().wait()  # 永久运行


def _start_in_thread(cert_file, key_file):
    def _run():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(_run_server(cert_file, key_file))
        except OSError as e:
            if e.errno == 98 or 'address already in use' in str(e).lower():
                print(f'[HTTPS Bridge] ⚠️  端口 {BRIDGE_PORT} 已被占用，桥接未启动。')
                print(f'               可通过环境变量更换端口: COMFY_BRIDGE_PORT=9998')
            else:
                print(f'[HTTPS Bridge] 启动失败: {e}')
                traceback.print_exc()

    thread = threading.Thread(target=_run, daemon=True, name='https-bridge')
    thread.start()

# ── ComfyUI Node 注册 ─────────────────────────────────────────────────────────

class HttpsBridgeNode:
    """显示桥接状态的信息节点（可选使用，桥接服务在后台自动运行）。"""

    @classmethod
    def INPUT_TYPES(cls):
        return {'required': {}}

    RETURN_TYPES    = ('STRING',)
    FUNCTION        = 'get_status'
    CATEGORY        = '🔒 HTTPS Bridge'
    OUTPUT_NODE     = True

    def get_status(self):
        return (f'HTTPS Bridge: https://localhost:{BRIDGE_PORT}',)


NODE_CLASS_MAPPINGS        = {'HttpsBridgeNode': HttpsBridgeNode}
NODE_DISPLAY_NAME_MAPPINGS = {'HttpsBridgeNode': '🔒 HTTPS Bridge Info'}

# ── 入口：ComfyUI 加载时自动执行 ─────────────────────────────────────────────

cert_file, key_file = generate_cert()
if cert_file:
    _start_in_thread(cert_file, key_file)