import sys
import argparse
import platform
import ctypes
import winreg
from mitmproxy.tools.main import mitmdump
from proxy_addon import load_config


class SystemProxyManager:
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.original_enable = None
        self.original_server = None

    def _apply_windows_internet_options(self):
        option_refresh = 37
        option_settings_changed = 39
        internet_set_option = ctypes.windll.Wininet.InternetSetOptionW
        internet_set_option(0, option_settings_changed, 0, 0)
        internet_set_option(0, option_refresh, 0, 0)

    def enable(self):
        if platform.system().lower() != "windows":
            return
        key_path = r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, key_path, 0, winreg.KEY_READ | winreg.KEY_WRITE)
        try:
            self.original_enable, _ = winreg.QueryValueEx(key, "ProxyEnable")
        except FileNotFoundError:
            self.original_enable = 0
        try:
            self.original_server, _ = winreg.QueryValueEx(key, "ProxyServer")
        except FileNotFoundError:
            self.original_server = ""
        winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, 1)
        winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, f"{self.host}:{self.port}")
        winreg.CloseKey(key)
        self._apply_windows_internet_options()

    def disable(self):
        if platform.system().lower() != "windows":
            return
        if self.original_enable is None or self.original_server is None:
            return
        key_path = r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, key_path, 0, winreg.KEY_READ | winreg.KEY_WRITE)
        winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, self.original_enable)
        winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, self.original_server)
        winreg.CloseKey(key)
        self._apply_windows_internet_options()


def start_proxy(config_path: str = "config.json", port: int = 8080, host: str = "127.0.0.1", enable_system_proxy: bool = True):
    config = load_config(config_path)

    proxy_config = config.get('proxy', {})
    listen_port = port or proxy_config.get('listen_port', 8080)
    listen_host = host or proxy_config.get('listen_host', '127.0.0.1')

    print(f"\n{'='*60}")
    print(f"LLM Proxy Server")
    print(f"{'='*60}")
    print(f"Listening on: {listen_host}:{listen_port}")
    print(f"Config file: {config_path}")
    print(f"Database: {config.get('database', {}).get('path', 'llm_data.db')}")
    if enable_system_proxy and platform.system().lower() == "windows":
        print("System proxy: enabled for current session")
    print(f"{'='*60}\n")

    manager = None
    if enable_system_proxy:
        manager = SystemProxyManager(listen_host, listen_port)
        manager.enable()

    sys.argv = [
        'mitmdump',
        '-q',
        '-s', 'proxy_addon.py',
        '--listen-host', listen_host,
        '--listen-port', str(listen_port),
        '--set', 'block_global=false',
        '--set', 'flow_detail=0'
    ]

    try:
        mitmdump()
    finally:
        if manager is not None:
            manager.disable()


def cli_main():
    parser = argparse.ArgumentParser(description='Start LLM Proxy Server')
    parser.add_argument('--config', '-c', type=str, default='config.json',
                       help='Path to config file')
    parser.add_argument('--port', '-p', type=int, default=None,
                       help='Listen port (overrides config)')
    parser.add_argument('--host', '-H', type=str, default=None,
                       help='Listen host (overrides config)')
    parser.add_argument('--no-system-proxy', action='store_true',
                       help='Do not modify system proxy settings')

    args = parser.parse_args()

    start_proxy(args.config, args.port, args.host, not args.no_system_proxy)


if __name__ == '__main__':
    cli_main()