"""
FreeLLMAPI Catalog Proxy
Transparent caching reverse-proxy for FreeLLMAPI model catalog sync.
Fetches signed catalog from https://api.freellmapi.co/v1/latest via SOCKS5 proxy,
preserves the official Ed25519 digital signature (x-catalog-signature),
and serves it locally to the FreeLLMAPI container.
"""

import os
import sys
import time
import json
import socket
import ssl
from http.server import HTTPServer, BaseHTTPRequestHandler

CACHE_DIR = os.environ.get("CACHE_DIR", "/data")
CONFIG_FILE = os.path.join(CACHE_DIR, "proxy_config.json")

# Defaults
PORT = int(os.environ.get("PORT", "8080"))
PROXY_HOST = os.environ.get("PROXY_HOST", "77.239.127.49")
PROXY_PORT = int(os.environ.get("PROXY_PORT", "18322"))
PROXY_USER = os.environ.get("PROXY_USER", "")
PROXY_PASS = os.environ.get("PROXY_PASS", "")

if os.path.exists(CONFIG_FILE):
    try:
        with open(CONFIG_FILE, "r") as f:
            c = json.load(f)
            PROXY_HOST = c.get("proxy_host", PROXY_HOST)
            PROXY_PORT = int(c.get("proxy_port", PROXY_PORT))
            PROXY_USER = c.get("proxy_user", PROXY_USER)
            PROXY_PASS = c.get("proxy_pass", PROXY_PASS)
            print("[proxy] Loaded credentials from proxy_config.json")
    except Exception as e:
        print(f"[proxy] Error loading config: {e}")

CACHE_FILE = os.path.join(CACHE_DIR, "catalog_cache.json")
SIG_FILE = os.path.join(CACHE_DIR, "catalog_cache.sig")
META_FILE = os.path.join(CACHE_DIR, "catalog_cache.meta")
CACHE_TTL = int(os.environ.get("CACHE_TTL", "43200"))  # 12 hours


def fetch_upstream(path, timeout=25):
    s = socket.create_connection((PROXY_HOST, PROXY_PORT), timeout=timeout)
    s.sendall(b"\x05\x01\x02")
    res = s.recv(2)
    if res != b"\x05\x02":
        raise RuntimeError(f"SOCKS5 auth choice failed: {res}")
    auth_payload = (
        b"\x01"
        + bytes([len(PROXY_USER)])
        + PROXY_USER.encode()
        + bytes([len(PROXY_PASS)])
        + PROXY_PASS.encode()
    )
    s.sendall(auth_payload)
    res = s.recv(2)
    if res != b"\x01\x00":
        raise RuntimeError(f"SOCKS5 auth rejected: {res}")
    target_host = b"api.freellmapi.co"
    req = b"\x05\x01\x00\x03" + bytes([len(target_host)]) + target_host + (443).to_bytes(2, "big")
    s.sendall(req)
    res = s.recv(4)
    if res[1] != 0:
        raise RuntimeError(f"SOCKS5 connect error code: {res[1]}")
    if res[3] == 0x01:
        s.recv(6)
    elif res[3] == 0x03:
        s.recv(s.recv(1)[0] + 2)
    elif res[3] == 0x04:
        s.recv(18)

    ctx = ssl.create_default_context()
    ss = ctx.wrap_socket(s, server_hostname="api.freellmapi.co")
    http_req = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: api.freellmapi.co\r\n"
        f"User-Agent: freellmapi-catalog-proxy/1.0\r\n"
        f"Connection: close\r\n\r\n"
    ).encode()
    ss.sendall(http_req)

    chunks = []
    while True:
        data = ss.recv(65536)
        if not data:
            break
        chunks.append(data)
    ss.close()
    raw = b"".join(chunks)
    header_data, body = raw.split(b"\r\n\r\n", 1)
    status_line = header_data.split(b"\r\n")[0].decode("latin1")
    status_code = int(status_line.split()[1])
    headers = {}
    for line in header_data.decode("latin1").split("\r\n")[1:]:
        if ":" in line:
            k, v = line.split(":", 1)
            headers[k.strip().lower()] = v.strip()
    return status_code, headers, body


def get_catalog(force=False):
    now = time.time()
    if not force and os.path.exists(CACHE_FILE) and os.path.exists(SIG_FILE) and os.path.exists(META_FILE):
        try:
            with open(META_FILE, "r") as f:
                meta = json.load(f)
            if now - meta.get("timestamp", 0) < CACHE_TTL:
                with open(CACHE_FILE, "rb") as f:
                    body = f.read()
                with open(SIG_FILE, "r") as f:
                    sig = f.read().strip()
                return body, sig, meta.get("version", "")
        except Exception as e:
            print(f"[proxy] Cache read error: {e}")

    try:
        print("[proxy] Fetching fresh catalog from upstream...")
        code, hdrs, body = fetch_upstream("/v1/latest")
        if code == 200 and "x-catalog-signature" in hdrs and len(body) > 100000:
            sig = hdrs["x-catalog-signature"]
            try:
                parsed = json.loads(body)
                version = parsed.get("version", "")
            except Exception:
                version = ""
            os.makedirs(CACHE_DIR, exist_ok=True)
            with open(CACHE_FILE, "wb") as f:
                f.write(body)
            with open(SIG_FILE, "w") as f:
                f.write(sig)
            with open(META_FILE, "w") as f:
                json.dump({"timestamp": now, "version": version, "size": len(body)}, f)
            for path in [CACHE_FILE, SIG_FILE, META_FILE]:
                try:
                    os.chmod(path, 0o666)
                except Exception:
                    pass
            print(f"[proxy] Cached catalog v{version} ({len(body)} bytes)")
            return body, sig, version
        else:
            raise RuntimeError(f"Upstream returned status {code}, headers: {hdrs}")
    except Exception as e:
        print(f"[proxy] Fetch upstream failed: {e}")
        if os.path.exists(CACHE_FILE) and os.path.exists(SIG_FILE):
            print("[proxy] Falling back to existing cached catalog")
            with open(CACHE_FILE, "rb") as f:
                body = f.read()
            with open(SIG_FILE, "r") as f:
                sig = f.read().strip()
            return body, sig, "cached-fallback"
        raise


class Handler(BaseHTTPRequestHandler):
    def do_HEAD(self):
        if self.path.startswith("/v1/latest"):
            try:
                body, sig, ver = get_catalog()
                self.send_response(200)
                self.send_header("Content-Type", "application/json; charset=utf-8")
                self.send_header("x-catalog-signature", sig)
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
            except Exception as e:
                self.send_response(502)
                self.end_headers()
        else:
            self.send_response(404)
            self.end_headers()

    def do_GET(self):
        if self.path.startswith("/v1/latest"):
            try:
                body, sig, ver = get_catalog()
                self.send_response(200)
                self.send_header("Content-Type", "application/json; charset=utf-8")
                self.send_header("x-catalog-signature", sig)
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                print(f"[proxy] Served catalog v{ver} ({len(body)} bytes) to {self.client_address[0]}")
            except Exception as e:
                self.send_response(502)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": str(e)}).encode())
        elif self.path in ("/health", "/ping", "/"):
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            has_cache = os.path.exists(CACHE_FILE)
            self.wfile.write(json.dumps({"status": "ok", "cached": has_cache}).encode())
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        print(f"[proxy] {self.address_string()} - {format % args}")


if __name__ == "__main__":
    print(f"[proxy] Starting catalog-proxy on 0.0.0.0:{PORT}...")
    try:
        get_catalog()
    except Exception as e:
        print(f"[proxy] Initial catalog fetch warning: {e}")
    server = HTTPServer(("0.0.0.0", PORT), Handler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass