#!/usr/bin/env python3
"""Generate Hermes and OpenWebUI MCP registrations from one JSON registry."""

from __future__ import annotations

import argparse
import json
import os
import pathlib
import sqlite3
import time


BEGIN = "# BEGIN MANAGED MCP SERVERS"
END = "# END MANAGED MCP SERVERS"


def env_file(path: str) -> dict[str, str]:
    values: dict[str, str] = {}
    source = pathlib.Path(path)
    if not source.is_file():
        return values
    for raw in source.read_text(encoding="utf-8", errors="replace").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key.strip()] = value.strip().strip('"').strip("'")
    return values


def enabled(item: dict) -> bool:
    required = item.get("required_file")
    if required and not pathlib.Path(required).is_file():
        return False
    source = item.get("env_file")
    if source:
        values = env_file(source)
        url_ready = bool(item.get("url")) or bool(values.get(item.get("url_env", "")))
        key_ready = not item.get("key_env") or bool(values.get(item["key_env"]))
        return url_ready and key_ready
    return True


def resolved(item: dict) -> tuple[str, str]:
    if item.get("env_file"):
        values = env_file(item["env_file"])
        url = item.get("url") or values[item["url_env"]]
        key = values.get(item.get("key_env", ""), "")
        return url, key
    return item["url"], ""


def active(registry: pathlib.Path, client: str) -> list[dict]:
    document = json.loads(registry.read_text(encoding="utf-8"))
    if document.get("version") != 1 or not isinstance(document.get("servers"), list):
        raise SystemExit("Unsupported MCP registry schema")
    return [item for item in document["servers"] if client in item.get("clients", []) and enabled(item)]


def yaml_quote(value: str) -> str:
    return json.dumps(value, ensure_ascii=False)


def hermes_block(items: list[dict]) -> str:
    lines = [BEGIN, "mcp_servers:"]
    for item in items:
        url, key = resolved(item)
        lines.extend([
            f"  {item.get('hermes_id', item['id'])}:",
            f"    url: {yaml_quote(url)}",
        ])
        if key:
            lines.extend(["    headers:", f"      Authorization: {yaml_quote('Bearer ' + key)}"])
        lines.extend([
            f"    timeout: {int(item.get('timeout', 300))}",
            "    connect_timeout: 30",
            "    supports_parallel_tool_calls: false",
        ])
    lines.append(END)
    return "\n".join(lines) + "\n"


def update_hermes(path: pathlib.Path, block: str) -> None:
    if not path.is_file():
        return
    text = path.read_text(encoding="utf-8")
    if BEGIN in text and END in text:
        prefix, rest = text.split(BEGIN, 1)
        _, suffix = rest.split(END, 1)
        text = prefix.rstrip() + "\n\n" + block + suffix.lstrip("\n")
    else:
        marker = "\nmcp_servers:"
        if marker in text:
            text = text.split(marker, 1)[0].rstrip() + "\n\n" + block
        else:
            text = text.rstrip() + "\n\n" + block
    path.write_text(text, encoding="utf-8")


def openwebui_connection(item: dict) -> dict:
    url, key = resolved(item)
    config = {"enable": True, "access_grants": []}
    if item.get("functions"):
        config["function_name_filter_list"] = item["functions"]
    return {
        "url": url, "path": "", "type": "mcp",
        "auth_type": item.get("auth_type", "none"), "headers": None,
        "key": key, "config": config,
        "info": {"id": item["id"], "name": item["name"], "description": item["description"]},
    }


def update_openwebui(db: pathlib.Path, items: list[dict]) -> None:
    con = sqlite3.connect(db)
    now = int(time.time())
    row = con.execute("select value from config where key=?", ("tool_server.connections",)).fetchone()
    old = json.loads(row[0]) if row else []
    if not isinstance(old, list):
        raise SystemExit("Unexpected OpenWebUI tool_server.connections format")
    managed_ids = {
        "athena-platform", "athena-operator-local", "web-general-local", "github-local",
        "homeassistant-local", "arr-local", "navidrome-local", "deemix-local", "mua",
        "mua-readonly-local", "athena-terminal-local", "unraid-readonly-local", "web-local",
    }
    keep = [entry for entry in old if str((entry.get("info") or {}).get("id", "")) not in managed_ids]
    keep.extend(openwebui_connection(item) for item in items)
    with con:
        con.execute(
            """insert into config (key,value,updated_at) values (?,?,?)
               on conflict(key) do update set value=excluded.value,updated_at=excluded.updated_at""",
            ("tool_server.connections", json.dumps(keep, ensure_ascii=False), now),
        )
    con.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--registry", type=pathlib.Path, required=True)
    parser.add_argument("--hermes", type=pathlib.Path, action="append", default=[])
    parser.add_argument("--openwebui-db", type=pathlib.Path)
    args = parser.parse_args()
    if args.hermes:
        block = hermes_block(active(args.registry, "hermes"))
        for path in args.hermes:
            update_hermes(path, block)
    if args.openwebui_db:
        update_openwebui(args.openwebui_db, active(args.registry, "openwebui"))
    print("MCP_CLIENT_SYNC_OK")


if __name__ == "__main__":
    main()