#!/usr/bin/env python3
from __future__ import annotations

import json
import sys
from pathlib import Path


STATE_FILE = Path("/etc/caddy/managed/desineuron-routes.json")
SNIPPET_FILE = Path("/etc/caddy/managed/desineuron-routes.caddy")


def load_routes() -> dict[str, dict]:
    if STATE_FILE.exists():
        return json.loads(STATE_FILE.read_text(encoding="utf-8"))
    return {}


def save_routes(routes: dict[str, dict]) -> None:
    STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    STATE_FILE.write_text(json.dumps(routes, indent=2), encoding="utf-8")


def render_routes(routes: dict[str, dict]) -> None:
    lines: list[str] = []
    for hostname, route in sorted(routes.items()):
        lines.extend(
            [
                f"{hostname} {{",
                "\ttls /etc/caddy/tls/fullchain.pem /etc/caddy/tls/privkey.pem",
                "\tlog {",
                "\t\toutput file /var/log/caddy/access.log",
                "\t\tformat json",
                "\t}",
                f"\treverse_proxy {route['scheme']}://{route['target_host']}:{route['target_port']} {{",
                "\t\theader_up Host {host}",
                "\t\theader_up X-Forwarded-Host {host}",
                "\t\theader_up X-Forwarded-Proto {scheme}",
                "\t\theader_up X-Forwarded-For {remote_host}",
                "\t}",
                "}",
                "",
            ]
        )
    SNIPPET_FILE.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")

    # Generate a dedicated upstream include exclusively for velocity.desineuron.in/llm
    llm_inc = Path("/etc/caddy/managed/llm_upstream.caddy_inc")
    if "llm.desineuron.in" in routes:
        route = routes["llm.desineuron.in"]
        llm_inc.write_text(
            f"handle_path /llm/* {{\n"
            f"\treverse_proxy {route['scheme']}://{route['target_host']}:{route['target_port']} {{\n"
            f"\t\theader_up Host {{host}}\n"
            f"\t\theader_up X-Forwarded-For {{remote_host}}\n"
            f"\t\tflush_interval -1\n"
            f"\t\theader_down X-Accel-Buffering no\n"
            f"\t}}\n"
            f"}}\n",
            encoding="utf-8",
        )
    else:
        llm_inc.write_text("", encoding="utf-8")


def main() -> int:
    if len(sys.argv) < 2:
        print("usage: manage_desineuron_routes.py <upsert|delete|list> [payload|hostname]")
        return 1
    command = sys.argv[1]
    routes = load_routes()
    if command == "upsert":
        payload = json.loads(sys.argv[2])
        routes[payload["hostname"]] = payload
        save_routes(routes)
        render_routes(routes)
        print(json.dumps({"status": "ok", "action": "upsert", "hostname": payload["hostname"]}))
        return 0
    if command == "delete":
        hostname = sys.argv[2]
        routes.pop(hostname, None)
        save_routes(routes)
        render_routes(routes)
        print(json.dumps({"status": "ok", "action": "delete", "hostname": hostname}))
        return 0
    if command == "list":
        print(json.dumps(routes, indent=2))
        return 0
    print(f"unknown command: {command}")
    return 1


if __name__ == "__main__":
    raise SystemExit(main())