#!/usr/bin/env python3
"""Schlüssel-Proxy für die KI-API von ITM (API-Werk) – nur für die Entwicklung.

Liest beim Start die API-Schlüssel aus der Windows-Anmeldeinformationsverwaltung,
hält sie nur im Arbeitsspeicher und setzt sie in weitergeleitete Anfragen ein.
Aufrufer (Claude Code, Tests, die App in Sail) schicken Anfragen ohne Schlüssel.

Der Proxy wird ausschließlich vom Entwickler in einem eigenen Terminal gestartet,
nie von Claude Code. Er gibt Schlüssel nie aus und protokolliert keine Header.
"""

from __future__ import annotations

import argparse
import base64
import http.client
import json
import logging
import shutil
import ssl
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Callable, Optional
from urllib.parse import urlsplit

# Pfadpräfix → Eintrag in der Windows-Anmeldeinformationsverwaltung
ROUTES: dict[str, str] = {
    "/v1/ki/": "ki-bauin/itm-ki",
    "/v1/dataloader/": "ki-bauin/dataloader",
    "/v1/ocr/": "ki-bauin/ocr",
}

DEFAULT_UPSTREAM = "https://api-werk.de"
DEFAULT_PORT = 8787
DEFAULT_LISTEN = ("127.0.0.1", "172.17.0.1")  # WSL selbst und die Docker-Brücke (Sail)
MAX_BODY_BYTES = 60 * 1024 * 1024  # Dateigrenze der API ist 50 MB
UPSTREAM_TIMEOUT = 300

HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailer", "trailers", "transfer-encoding", "upgrade",
}

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

# Liest generische Einträge über die Windows-API CredRead. Ausgabe je Eintrag:
# "<eintrag>\t<base64 des Werts in UTF-8>" oder "<eintrag>\t-", wenn er fehlt.
_POWERSHELL_TEMPLATE = r"""
$ErrorActionPreference = 'Stop'
Add-Type -TypeDefinition @'
using System;
using System.Runtime.InteropServices;
public static class KiBauinTresor {
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    private struct CREDENTIAL {
        public int Flags; public int Type; public string TargetName; public string Comment;
        public System.Runtime.InteropServices.ComTypes.FILETIME LastWritten;
        public int CredentialBlobSize; public IntPtr CredentialBlob; public int Persist;
        public int AttributeCount; public IntPtr Attributes; public string TargetAlias; public string UserName;
    }
    [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
    private static extern bool CredRead(string target, int type, int flags, out IntPtr credential);
    [DllImport("advapi32.dll")]
    private static extern void CredFree(IntPtr credential);
    public static string Read(string target) {
        IntPtr p;
        if (!CredRead(target, 1, 0, out p)) { return null; }
        try {
            CREDENTIAL c = (CREDENTIAL)Marshal.PtrToStructure(p, typeof(CREDENTIAL));
            if (c.CredentialBlobSize == 0) { return ""; }
            return Marshal.PtrToStringUni(c.CredentialBlob, c.CredentialBlobSize / 2);
        } finally { CredFree(p); }
    }
}
'@
foreach ($t in @(__TARGETS__)) {
    $v = [KiBauinTresor]::Read($t)
    if ($v -eq $null) { [Console]::Out.WriteLine($t + "`t-") }
    else { [Console]::Out.WriteLine($t + "`t" + [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($v.Trim()))) }
}
"""


def parse_tresor_output(output: str) -> dict[str, Optional[str]]:
    """Wertet die Ausgabe des PowerShell-Skripts aus."""
    result: dict[str, Optional[str]] = {}
    for line in output.splitlines():
        line = line.strip()
        if not line or "\t" not in line:
            continue
        target, value = line.split("\t", 1)
        result[target] = None if value == "-" else base64.b64decode(value).decode("utf-8")
    return result


def read_windows_credentials(targets: list[str]) -> dict[str, Optional[str]]:
    """Liest die angegebenen Einträge aus der Windows-Anmeldeinformationsverwaltung."""
    powershell = shutil.which("powershell.exe")
    if powershell is None:
        raise RuntimeError("powershell.exe nicht gefunden – läuft das unter WSL mit Windows-Interop?")
    quoted = ", ".join("'" + t.replace("'", "''") + "'" for t in targets)
    script = _POWERSHELL_TEMPLATE.replace("__TARGETS__", quoted)
    encoded = base64.b64encode(script.encode("utf-16-le")).decode("ascii")
    completed = subprocess.run(
        [powershell, "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-EncodedCommand", encoded],
        capture_output=True,
        timeout=60,
    )
    if completed.returncode != 0:
        # Die Meldung von PowerShell enthält keine Schlüssel; sie kommt in der Windows-Codepage.
        detail = completed.stderr.decode("cp850", errors="replace").strip().splitlines()[:3]
        raise RuntimeError("Lesen aus dem Windows-Tresor fehlgeschlagen: " + " | ".join(detail))
    found = parse_tresor_output(completed.stdout.decode("ascii", errors="replace"))
    return {t: found.get(t) for t in targets}


class _LimitedReader:
    """Liest genau `remaining` Bytes aus dem Eingangsstrom, damit der Request-Body gestreamt wird."""

    def __init__(self, stream, length: int):
        self._stream = stream
        self._remaining = length

    def read(self, size: int = -1) -> bytes:
        if self._remaining <= 0:
            return b""
        if size < 0 or size > self._remaining:
            size = self._remaining
        chunk = self._stream.read(size)
        self._remaining -= len(chunk)
        return chunk


class ProxyHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.0"  # Antworten enden mit dem Schließen der Verbindung; einfach für Streaming
    server_version = "ki-proxy"
    server: "KiProxyServer"

    def do_GET(self) -> None:
        self._forward()

    def do_POST(self) -> None:
        self._forward()

    def do_PUT(self) -> None:
        self._forward()

    def do_PATCH(self) -> None:
        self._forward()

    def do_DELETE(self) -> None:
        self._forward()

    def log_message(self, format: str, *args) -> None:  # noqa: A002 – Standardausgabe abschalten
        return

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _forward(self) -> None:
        started = time.monotonic()
        path_only = self.path.split("?", 1)[0]
        target = next((t for prefix, t in ROUTES.items() if path_only.startswith(prefix)), None)
        status = 0
        sent = 0
        try:
            if target is None:
                status = 404
                self._send_json(404, {"error": "Unbekannter Pfad. Erlaubt: " + ", ".join(ROUTES)})
                return
            key = self.server.keys.get(target)
            if not key:
                status = 503
                self._send_json(503, {"error": f"Schlüssel '{target}' fehlt im Windows-Tresor. Proxy nach dem Eintragen neu starten."})
                return
            if self.headers.get("Transfer-Encoding", "").lower() == "chunked":
                status = 411
                self._send_json(411, {"error": "Request-Body mit Content-Length senden, nicht chunked."})
                return
            length = int(self.headers.get("Content-Length") or 0)
            if length > MAX_BODY_BYTES:
                status = 413
                self._send_json(413, {"error": "Request zu groß für den Proxy."})
                return

            headers = {
                name: value
                for name, value in self.headers.items()
                if name.lower() not in HOP_BY_HOP | {"authorization", "host", "content-length"}
            }
            headers["Authorization"] = f"Bearer {key}"
            if length or self.command in ("POST", "PUT", "PATCH"):
                headers["Content-Length"] = str(length)
            body = _LimitedReader(self.rfile, length) if length else None

            conn = self.server.connect_upstream()
            try:
                conn.request(self.command, self.server.upstream_base_path + self.path, body=body, headers=headers)
                response = conn.getresponse()
                status = response.status
                self.send_response(response.status, response.reason)
                for name, value in response.getheaders():
                    if name.lower() not in HOP_BY_HOP | {"server", "date"}:
                        self.send_header(name, value)
                self.end_headers()
                while True:
                    chunk = response.read1(65536)
                    if not chunk:
                        break
                    self.wfile.write(chunk)
                    self.wfile.flush()
                    sent += len(chunk)
            finally:
                conn.close()
        except (OSError, http.client.HTTPException) as exc:
            if status == 0:
                status = 502
                try:
                    self._send_json(502, {"error": f"KI-API nicht erreichbar ({type(exc).__name__})."})
                except OSError:
                    pass
        finally:
            log.info("%s %s → %s, %d Bytes, %.0f ms", self.command, path_only, status or "-", sent, (time.monotonic() - started) * 1000)


class KiProxyServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, address: tuple[str, int], keys: dict[str, Optional[str]], upstream: str):
        super().__init__(address, ProxyHandler)
        self.keys = keys
        parts = urlsplit(upstream)
        self.upstream_scheme = parts.scheme
        self.upstream_host = parts.hostname or ""
        self.upstream_port = parts.port or (443 if parts.scheme == "https" else 80)
        self.upstream_base_path = parts.path.rstrip("/")
        self._ssl_context = ssl.create_default_context() if parts.scheme == "https" else None

    def connect_upstream(self) -> http.client.HTTPConnection:
        if self.upstream_scheme == "https":
            return http.client.HTTPSConnection(self.upstream_host, self.upstream_port, timeout=UPSTREAM_TIMEOUT, context=self._ssl_context)
        return http.client.HTTPConnection(self.upstream_host, self.upstream_port, timeout=UPSTREAM_TIMEOUT)


def start_servers(hosts: list[str], port: int, keys: dict[str, Optional[str]], upstream: str) -> list[KiProxyServer]:
    servers = []
    for host in hosts:
        try:
            server = KiProxyServer((host, port), keys, upstream)
        except OSError as exc:
            log.warning("Kann nicht auf %s:%d lauschen (%s) – übersprungen.", host, port, exc.strerror)
            continue
        threading.Thread(target=server.serve_forever, daemon=True).start()
        servers.append(server)
    return servers


def main(argv: Optional[list[str]] = None, reader: Callable[[list[str]], dict[str, Optional[str]]] = read_windows_credentials) -> int:
    parser = argparse.ArgumentParser(description="Schlüssel-Proxy für die KI-API von ITM (nur Entwicklung).")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument("--listen", default=",".join(DEFAULT_LISTEN), help="Adressen, durch Komma getrennt")
    parser.add_argument("--upstream", default=DEFAULT_UPSTREAM)
    parser.add_argument("--check", action="store_true", help="Nur prüfen, welche Schlüssel im Tresor stehen, dann beenden")
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S")

    keys = reader(list(ROUTES.values()))
    for target in ROUTES.values():
        log.info("%-22s %s", target, "gefunden" if keys.get(target) else "FEHLT")
    if args.check:
        return 0 if all(keys.values()) else 1

    hosts = [h.strip() for h in args.listen.split(",") if h.strip()]
    servers = start_servers(hosts, args.port, keys, args.upstream)
    if not servers:
        log.error("Keine Adresse verfügbar, Proxy beendet.")
        return 1
    log.info("Proxy läuft auf %s → %s (Strg+C beendet)", ", ".join(f"{s.server_address[0]}:{s.server_address[1]}" for s in servers), args.upstream)
    try:
        while True:
            time.sleep(3600)
    except KeyboardInterrupt:
        log.info("Proxy beendet.")
    finally:
        for server in servers:
            server.shutdown()
    return 0


if __name__ == "__main__":
    sys.exit(main())