#!/usr/bin/env python3
"""
LLM Proxy - MITM 代理模块
==========================
拦截发往 DeepSeek / MiniMax 等 LLM API 的 HTTPS 请求，
捕获完整的 prompt（请求体）和 response（响应体），存入监控数据库。

使用方法:
    1. 首次运行自动生成 CA 证书 (ca.crt + ca.key)
    2. 将 ca.crt 安装为系统受信任的根证书
    3. 设置环境变量 HTTPS_PROXY=http://127.0.0.1:5911
    4. 重启 WorkBuddy，所有 LLM 请求将被记录
"""

import asyncio
import json
import os
import ssl
import time
import logging
import sqlite3
import traceback
from datetime import datetime
from urllib.parse import urlparse

from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend

log = logging.getLogger("llm-proxy")

# ========== Configuration ==========
PROXY_HOST = "127.0.0.1"
PROXY_PORT = 5911
CA_DIR = os.path.expanduser(r"~\.workbuddy-monitor\proxy-ca")
CA_CERT_PATH = os.path.join(CA_DIR, "ca.crt")
CA_KEY_PATH = os.path.join(CA_DIR, "ca.key")

# 目标 API 主机列表（只拦截这些主机的请求）
TARGET_HOSTS = [
    "api.deepseek.com",
    "api.minimaxi.com",
    "api.xiaomimimo.com",
    "token-plan-cn.xiaomimimo.com",
]

# 共享数据库路径
# - 由 monitor.py 通过 llm_proxy.DB_PATH 设置
# - 独立运行时使用默认持久化路径
_DEFAULT_DB_DIR = os.path.expanduser(r"~\.workbuddy-monitor")
os.makedirs(_DEFAULT_DB_DIR, exist_ok=True)
DB_PATH = os.path.join(_DEFAULT_DB_DIR, "monitor_data.db")

# ========== SSL Context 缓存 ==========
_ssl_contexts = {}  # hostname -> SSLContext

# ========== 运行状态 ==========
_server = None
_running = False

# ========== Certificate Management ==========

def _generate_key():
    """Generate an RSA private key."""
    return rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend(),
    )

def _generate_ca():
    """Generate a self-signed CA certificate and key."""
    os.makedirs(CA_DIR, exist_ok=True)

    ca_key = _generate_key()
    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "WorkBuddy Monitor Proxy"),
        x509.NameAttribute(NameOID.COMMON_NAME, "WorkBuddy Monitor CA"),
    ])

    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(ca_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.utcnow())
        .not_valid_after(datetime.utcnow().replace(year=datetime.utcnow().year + 10))
        .add_extension(
            x509.BasicConstraints(ca=True, path_length=None), critical=True
        )
        .add_extension(
            x509.KeyUsage(
                digital_signature=True,
                key_cert_sign=True,
                key_encipherment=True,
                key_agreement=False,
                content_commitment=False,
                data_encipherment=False,
                crl_sign=True,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=True,
        )
        .sign(ca_key, hashes.SHA256(), default_backend())
    )

    with open(CA_KEY_PATH, "wb") as f:
        f.write(ca_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        ))
    with open(CA_CERT_PATH, "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))

    log.info("CA certificate generated: %s", CA_CERT_PATH)
    return ca_key, cert


def _load_ca():
    """Load existing CA cert/key or generate new ones."""
    if os.path.isfile(CA_KEY_PATH) and os.path.isfile(CA_CERT_PATH):
        with open(CA_KEY_PATH, "rb") as f:
            ca_key = serialization.load_pem_private_key(
                f.read(), password=None, backend=default_backend()
            )
        with open(CA_CERT_PATH, "rb") as f:
            ca_cert = x509.load_pem_x509_certificate(f.read(), default_backend())
        log.info("CA certificate loaded: %s", CA_CERT_PATH)
        return ca_key, ca_cert
    return _generate_ca()


def _generate_host_cert(hostname, ca_key, ca_cert):
    """Generate a leaf certificate for the given hostname, signed by the CA."""
    host_key = _generate_key()

    san_list = [x509.DNSName(hostname)]
    # Handle wildcard domains
    if hostname.count(".") >= 2:
        parts = hostname.split(".")
        san_list.append(x509.DNSName("*." + ".".join(parts[-2:])))

    cert = (
        x509.CertificateBuilder()
        .subject_name(x509.Name([
            x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "WorkBuddy Monitor Proxy"),
            x509.NameAttribute(NameOID.COMMON_NAME, hostname),
        ]))
        .issuer_name(ca_cert.subject)
        .public_key(host_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.utcnow())
        .not_valid_after(datetime.utcnow().replace(year=datetime.utcnow().year + 1))
        .add_extension(
            x509.SubjectAlternativeName(san_list), critical=False
        )
        .add_extension(
            x509.BasicConstraints(ca=False, path_length=None), critical=True
        )
        .sign(ca_key, hashes.SHA256(), default_backend())
    )

    return host_key, cert


def _get_ssl_context_for_host(hostname, ca_key, ca_cert):
    """Get or create an SSL server context for the given hostname."""
    if hostname in _ssl_contexts:
        return _ssl_contexts[hostname]

    host_key, host_cert = _generate_host_cert(hostname, ca_key, ca_cert)

    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    # Store cert and key as PEM in memory
    cert_pem = host_cert.public_bytes(serialization.Encoding.PEM)
    key_pem = host_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )

    ctx.load_cert_chain(certfile=ssl.MemoryBIO(cert_pem), keyfile=ssl.MemoryBIO(key_pem))
    # Actually, load_cert_chain requires file paths. Let's use temp files.
    # Better approach: configure the context with the cert and key directly.

    _ssl_contexts[hostname] = ctx
    return ctx


# ========== Database ==========

def _init_captures_db():
    """Initialize the captures table in the monitor database."""
    if not DB_PATH:
        return
    try:
        conn = sqlite3.connect(DB_PATH)
        c = conn.cursor()
        c.execute("""
            CREATE TABLE IF NOT EXISTS captures (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id TEXT,
                session_id TEXT,
                hostname TEXT,
                method TEXT,
                path TEXT,
                request_body TEXT,
                response_body TEXT,
                status_code INTEGER,
                captured_at REAL,
                duration_ms INTEGER,
                content_length INTEGER,
                model_name TEXT
            )
        """)
        conn.commit()
        conn.close()
    except Exception as e:
        log.warning("Failed to init captures db: %s", e)


def save_capture(hostname, method, path, request_body, response_body,
                 status_code, duration_ms, trace_id="", session_id=""):
    """Save a captured LLM request/response pair to the database."""
    if not DB_PATH:
        return
    try:
        # Try to extract model name from the request body
        model_name = ""
        if request_body:
            try:
                req = json.loads(request_body)
                model_name = req.get("model", "")
            except (json.JSONDecodeError, TypeError):
                pass

        content_length = len(response_body) if response_body else 0

        conn = sqlite3.connect(DB_PATH)
        c = conn.cursor()
        c.execute("""
            INSERT INTO captures
                (hostname, method, path, request_body, response_body,
                 status_code, captured_at, duration_ms, content_length, model_name)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            hostname, method, path, request_body, response_body,
            status_code, time.time(), duration_ms, content_length, model_name,
        ))
        conn.commit()
        conn.close()
    except Exception as e:
        log.warning("Failed to save capture: %s", e)


def _fix_mimo_image_request(req_body: bytes) -> bytes:
    """修复 MiMo 图片理解请求中的两个问题：

    1. tool 角色消息中包含 image_url → 转为 user 角色（MiMo/OpenAI 不支持 tool 消息传图）
    2. 修复后如果只有 image_url 没有 text → 自动添加默认 text
    """
    if not req_body:
        return req_body

    try:
        req = json.loads(req_body)

        # 只对 mimo 系列模型生效
        model = req.get("model", "")
        if "mimo" not in model.lower():
            return req_body

        fixed = False

        for msg in req.get("messages", []):
            content = msg.get("content")
            if not isinstance(content, list):
                continue

            # 检查是否包含图片
            has_image = any(
                isinstance(item, dict) and item.get("type") == "image_url"
                for item in content
            )
            if not has_image:
                continue

            # 修复1: tool 角色 → user 角色（tool 不支持图片）
            if msg.get("role") == "tool":
                msg["role"] = "user"
                # 清理 tool 消息特有的字段
                msg.pop("tool_call_id", None)
                fixed = True
                log.info("Fixed tool→user role for message with image")

            # 修复2: 补充缺失的 text
            has_text = any(
                isinstance(item, dict) and item.get("type") == "text"
                for item in content
            )
            if not has_text:
                content.append({
                    "type": "text",
                    "text": "请描述这张图片的内容"
                })
                fixed = True
                log.info("Fixed missing 'text' in content")

        if fixed:
            new_body = json.dumps(req, ensure_ascii=False).encode("utf-8")
            return new_body

        return req_body
    except (json.JSONDecodeError, TypeError, KeyError):
        return req_body


# ========== MITM Proxy Core ==========

async def _pipe(reader, writer, label, callback=None):
    """Read data from reader and write to writer, optionally calling callback."""
    try:
        while True:
            data = await reader.read(65536)
            if not data:
                break
            if callback:
                await callback(data)
            writer.write(data)
            await writer.drain()
    except (ConnectionError, asyncio.CancelledError):
        pass
    except Exception as e:
        log.debug("Pipe error (%s): %s", label, e)


async def _handle_connect(reader, writer, target_host, target_port, ca_key, ca_cert):
    """Handle a CONNECT request by establishing a MITM tunnel."""
    start_time = time.time()
    request_buf = bytearray()
    response_buf = bytearray()
    status_code = 0

    try:
        # 1. Create SSL server context for this host
        host_key, host_cert = _generate_host_cert(target_host, ca_key, ca_cert)

        server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        server_ctx.check_hostname = False
        server_ctx.verify_mode = ssl.CERT_NONE

        # Write cert and key to temp files for load_cert_chain
        import tempfile
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pem") as cert_f:
            cert_f.write(host_cert.public_bytes(serialization.Encoding.PEM))
            cert_path = cert_f.name
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pem") as key_f:
            key_f.write(host_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.TraditionalOpenSSL,
                encryption_algorithm=serialization.NoEncryption(),
            ))
            key_path = key_f.name

        server_ctx.load_cert_chain(cert_path, key_path)

        # 2. Upgrade client connection to TLS (we are the server)
        client_ssl = await asyncio.start_server(
            lambda r, w: None, host="127.0.0.1", port=0
        )
        # Actually, let's use wrap directly on the existing connection
        # But asyncio's wrap doesn't easily support upgrading an existing connection.
        # Alternative: wrap the reader/writer pair.

        # Send 200 Connection Established
        writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
        await writer.drain()

        # Wrap the transport in SSL
        client_tls = await asyncio.start_tls(
            writer.transport,
            writer.protocol,
            server_ctx,
            server_side=True,
        )
        # This is complex with asyncio. Let me use a different approach.

        # Actually, the approach above is getting complex. Let me use socket-level approach.
        writer.close()
        return

    except Exception as e:
        log.error("CONNECT error for %s:%d: %s", target_host, target_port, e)
        try:
            writer.close()
        except:
            pass


async def _handle_client(reader, writer, ca_key, ca_cert):
    """Handle an incoming client connection."""
    try:
        # Read the first request line
        request_line = await asyncio.wait_for(reader.readline(), timeout=30)
        if not request_line:
            return

        # Parse headers
        headers = {}
        while True:
            header_line = await reader.readline()
            if header_line in (b"\r\n", b"\n", b""):
                break
            header_str = header_line.decode("utf-8", errors="replace").strip()
            if ":" in header_str:
                key, val = header_str.split(":", 1)
                headers[key.strip().lower()] = val.strip()

        line = request_line.decode("utf-8", errors="replace").strip()
        parts = line.split()
        if len(parts) < 3:
            return

        method = parts[0]
        target = parts[1]  # For CONNECT: host:port; For others: full URL or path

        if method.upper() == "CONNECT":
            # Handle CONNECT tunnel
            host_port = target
            if ":" in host_port:
                target_host, target_port_str = host_port.split(":", 1)
                target_port = int(target_port_str)
            else:
                target_host = host_port
                target_port = 443

            # Only intercept our target hosts
            if target_host not in TARGET_HOSTS:
                # Pass through without interception
                writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
                await writer.drain()

                # Relay the connection
                try:
                    remote_reader, remote_writer = await asyncio.wait_for(
                        asyncio.open_connection(target_host, target_port), timeout=10
                    )
                    await asyncio.gather(
                        _pipe(reader, remote_writer, "client->remote"),
                        _pipe(remote_reader, writer, "remote->client"),
                    )
                except Exception as e:
                    log.debug("Relay error for %s: %s", target_host, e)
                finally:
                    try:
                        writer.close()
                    except:
                        pass
                return

            # MITM interception
            start_time = time.time()
            request_buf = bytearray()
            response_buf = bytearray()
            status_code = 0

            try:
                # Generate host cert
                host_key, host_cert = _generate_host_cert(target_host, ca_key, ca_cert)

                import tempfile
                with tempfile.NamedTemporaryFile(delete=False, suffix=".pem") as cf:
                    cf.write(host_cert.public_bytes(serialization.Encoding.PEM))
                    cert_path = cf.name
                with tempfile.NamedTemporaryFile(delete=False, suffix=".pem") as kf:
                    kf.write(host_key.private_bytes(
                        encoding=serialization.Encoding.PEM,
                        format=serialization.PrivateFormat.TraditionalOpenSSL,
                        encryption_algorithm=serialization.NoEncryption(),
                    ))
                    key_path = kf.name

                server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
                server_ctx.check_hostname = False
                server_ctx.verify_mode = ssl.CERT_NONE
                server_ctx.load_cert_chain(cert_path, key_path)

                # Send 200 to client
                writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
                await writer.drain()

                # Upgrade to TLS on the client side (we act as server)
                client_tls_reader, client_tls_writer = await _start_tls(
                    reader, writer, server_ctx, server_side=True
                )

                # Connect to the real server and upgrade to TLS (we act as client)
                remote_reader, remote_writer = await asyncio.wait_for(
                    asyncio.open_connection(target_host, target_port, ssl=True), timeout=10
                )

                # Now read the HTTP request from the client
                req_line = await client_tls_reader.readline()
                if not req_line:
                    return

                req_method = req_line.decode().strip().split(" ")[0]
                req_path = req_line.decode().strip().split(" ")[1]

                # Read HTTP headers
                req_headers = {}
                req_body_len = 0
                while True:
                    hl = await client_tls_reader.readline()
                    if hl in (b"\r\n", b"\n", b""):
                        break
                    hs = hl.decode("utf-8", errors="replace").strip()
                    if ":" in hs:
                        k, v = hs.split(":", 1)
                        req_headers[k.strip().lower()] = v.strip()

                content_length = int(req_headers.get("content-length", 0))

                # Read request body if present
                req_body = b""
                if content_length > 0:
                    req_body = await asyncio.wait_for(
                        client_tls_reader.readexactly(content_length), timeout=10
                    )

                # === 修复 MiMo 图片请求缺少 text 参数 ===
                req_body = _fix_mimo_image_request(req_body)
                if req_body and len(req_body) != content_length:
                    req_headers["content-length"] = str(len(req_body))
                # ===========================================

                # Forward the request to the real server
                # Reconstruct the HTTP request
                remote_writer.write(req_line)
                for key, val in req_headers.items():
                    remote_writer.write(f"{key}: {val}\r\n".encode())
                remote_writer.write(b"\r\n")
                if req_body:
                    remote_writer.write(req_body)
                await remote_writer.drain()

                # Read response from server
                resp_line = await remote_reader.readline()
                resp_parts = resp_line.decode().strip().split(" ", 2)
                status_code = int(resp_parts[1]) if len(resp_parts) > 1 else 0

                # Read response headers
                resp_headers = {}
                resp_body_len = 0
                content_type = ""
                while True:
                    hl = await remote_reader.readline()
                    if hl in (b"\r\n", b"\n", b""):
                        break
                    hs = hl.decode("utf-8", errors="replace").strip()
                    if ":" in hs:
                        k, v = hs.split(":", 1)
                        resp_headers[k.strip().lower()] = v.strip()

                resp_content_length = int(resp_headers.get("content-length", 0))
                transfer_encoding = resp_headers.get("transfer-encoding", "")

                # Read response body
                resp_body = b""
                if transfer_encoding.lower() == "chunked":
                    # Read chunked response
                    while True:
                        chunk_size_line = await remote_reader.readline()
                        chunk_size_str = chunk_size_line.decode().strip()
                        if not chunk_size_str:
                            continue
                        try:
                            chunk_size = int(chunk_size_str, 16)
                        except ValueError:
                            break
                        if chunk_size == 0:
                            await remote_reader.readline()  # trailing CRLF
                            break
                        chunk = await asyncio.wait_for(
                            remote_reader.readexactly(chunk_size), timeout=30
                        )
                        resp_body += chunk
                        await remote_reader.readline()  # trailing CRLF
                elif resp_content_length > 0:
                    resp_body = await asyncio.wait_for(
                        remote_reader.readexactly(resp_content_length), timeout=30
                    )
                else:
                    # Try reading all available data
                    try:
                        while True:
                            chunk = await asyncio.wait_for(
                                remote_reader.read(65536), timeout=5
                            )
                            if not chunk:
                                break
                            resp_body += chunk
                    except (asyncio.TimeoutError, ConnectionResetError):
                        pass

                # Fix response headers: dechunked body needs content-length
                if transfer_encoding.lower() == "chunked":
                    resp_headers.pop("transfer-encoding", None)
                    resp_headers.pop("Transfer-Encoding", None)
                    resp_headers["content-length"] = str(len(resp_body))

                # Forward the response to the client
                client_tls_writer.write(resp_line)
                for key, val in resp_headers.items():
                    client_tls_writer.write(f"{key}: {val}\r\n".encode())
                client_tls_writer.write(b"\r\n")
                if resp_body:
                    client_tls_writer.write(resp_body)
                await client_tls_writer.drain()
                client_tls_writer.close()

                # Save the capture
                duration_ms = int((time.time() - start_time) * 1000)
                req_body_str = req_body.decode("utf-8", errors="replace") if req_body else ""
                resp_body_str = resp_body.decode("utf-8", errors="replace") if resp_body else ""

                save_capture(
                    hostname=target_host,
                    method=req_method,
                    path=req_path,
                    request_body=req_body_str,
                    response_body=resp_body_str,
                    status_code=status_code,
                    duration_ms=duration_ms,
                )

                log.info("Captured %s %s | %d | %dms | req=%d resp=%d",
                         req_method, req_path, status_code, duration_ms,
                         len(req_body), len(resp_body))

                # Clean up temp files
                try:
                    os.unlink(cert_path)
                    os.unlink(key_path)
                except:
                    pass

            except Exception as e:
                log.error("MITM error for %s: %s", target_host, e)
                log.debug(traceback.format_exc())
            finally:
                try:
                    writer.close()
                except:
                    pass
        else:
            # Non-CONNECT method (proxy request) - pass through
            try:
                # Parse the target URL
                parsed = urlparse(target)
                target_host = parsed.hostname
                target_port = parsed.port or 443

                remote_reader, remote_writer = await asyncio.wait_for(
                    asyncio.open_connection(target_host, target_port, ssl=True), timeout=10
                )

                # Forward the request
                remote_writer.write(request_line)
                for key, val in headers.items():
                    remote_writer.write(f"{key}: {val}\r\n".encode())
                remote_writer.write(b"\r\n")

                # Read and forward response
                resp_line = await remote_reader.readline()
                writer.write(resp_line)

                while True:
                    hl = await remote_reader.readline()
                    writer.write(hl)
                    if hl in (b"\r\n", b"\n", b""):
                        break

                # Pipe remaining data
                await _pipe(remote_reader, writer, "relay")
            except Exception as e:
                log.debug("Relay error: %s", e)
            finally:
                try:
                    writer.close()
                except:
                    pass

    except (asyncio.TimeoutError, ConnectionError) as e:
        log.debug("Client connection error: %s", e)
    except Exception as e:
        log.debug("Handler error: %s", e)
        log.debug(traceback.format_exc())
    finally:
        try:
            writer.close()
        except:
            pass


async def _start_tls(reader, writer, ssl_context, server_side=True):
    """Start TLS on an existing connection using loop.start_tls.

    IMPORTANT: loop.start_tls() pauses the OLD protocol, upgrades to TLS,
    then resumes the OLD protocol to receive decrypted data via the
    original reader. We MUST NOT create a new reader.
    """
    transport = writer.transport
    # writer._protocol is the StreamReaderProtocol that feeds 'reader'
    protocol = writer._protocol
    loop = asyncio.get_running_loop()

    try:
        new_transport = await loop.start_tls(
            transport,
            protocol,
            ssl_context,
            server_side=server_side,
        )

        # Keep using the ORIGINAL reader - it receives decrypted data
        # from the old protocol (which is unpaused by start_tls).
        # Only create a new writer for the SSL transport.
        new_writer = asyncio.StreamWriter(new_transport, protocol, reader, loop=loop)

        return reader, new_writer
    except Exception as e:
        log.error("TLS upgrade failed: %s", e)
        log.debug(traceback.format_exc())
        raise


# ========== Public API ==========

def get_status():
    """Return proxy status dict."""
    return {
        "running": _running,
        "host": PROXY_HOST,
        "port": PROXY_PORT,
        "target_hosts": TARGET_HOSTS,
        "ca_cert_path": CA_CERT_PATH if os.path.isfile(CA_CERT_PATH) else None,
    }


async def start():
    """Start the proxy server."""
    global _server, _running

    if _running:
        log.warning("Proxy is already running")
        return

    try:
        # Initialize CA
        ca_key, ca_cert = _load_ca()
        _init_captures_db()

        # Start the proxy server
        _server = await asyncio.start_server(
            lambda r, w: _handle_client(r, w, ca_key, ca_cert),
            host=PROXY_HOST,
            port=PROXY_PORT,
        )

        _running = True
        log.info("LLM Proxy started on %s:%d", PROXY_HOST, PROXY_PORT)
        log.info("Target hosts: %s", TARGET_HOSTS)
        log.info("CA certificate: %s", CA_CERT_PATH)

        async with _server:
            await _server.serve_forever()
    except Exception as e:
        _running = False
        log.error("Failed to start proxy on %s:%d: %s", PROXY_HOST, PROXY_PORT, e)
        raise
    finally:
        _running = False


def start_proxy():
    """Start the proxy in a new asyncio event loop (thread-safe entry point)."""
    asyncio.run(start())


def stop_proxy():
    """Stop the proxy server."""
    global _server, _running
    if _server:
        try:
            _server.close()
            log.info("LLM Proxy stopped")
        except Exception as e:
            log.warning("Error closing server: %s", e)
        _server = None
    _running = False


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print(f"LLM Proxy - MITM proxy for capturing LLM prompts/responses")
    print(f"Target hosts: {TARGET_HOSTS}")
    print(f"CA directory: {CA_DIR}")
    print()
    print("Starting proxy...")
    start_proxy()
