#!/usr/bin/env python3
"""Fetch normalized model metadata from Plexus.

Plexus replaced its single-endpoint `/v1/models` (which used to include full
metadata) with a flattened alias list.  Enriched fields now require a
two-step lookup:

1. List aliases via the management API (`/v0/management/aliases`) using the
   admin key.  Each alias may carry a `metadata` block pointing at an
   external catalog entry (source + source_path).
2. For each alias with metadata, call the public `/v1/metadata/lookup`
   endpoint to retrieve the normalized record (context_length, capabilities,
   etc.).

Aliases without a `metadata` block (diff-apply, fix-json, embedding models,
etc.) are implicitly excluded.

Usage:
    plexus.py [--json | --tsv]
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, NotRequired, TypedDict

from models import Model, clean_name, emit_json, emit_tsv

PLEXUS_HOST = "harp-willow.snowy-egret.ts.net"
PLEXUS_PORT = 4000
BASE_URL = f"http://{PLEXUS_HOST}:{PLEXUS_PORT}"

OP_ITEM = "op://Shared/Plexus/password"

# Models that must never appear in any client config, even if Plexus reports
# them.  This is a safety net — they usually lack metadata and are filtered
# at the alias step anyway.
EXCLUDED_IDS: frozenset[str] = frozenset({
    "nomic-embed-text-v1.5",
    "diff-apply",
    "fix-json",
})


class AliasMetadata(TypedDict):
    source: str
    source_path: str
    overrides: NotRequired[dict[str, Any]]


class AliasConfig(TypedDict, total=False):
    target_groups: list[dict[str, Any]]
    metadata: AliasMetadata | None


def _get_admin_token() -> str:
    """Retrieve the Plexus admin key from 1Password.

    The `op` CLI must run without `OP_SERVICE_ACCOUNT_TOKEN` set, otherwise
    it tries to use the service account which lacks access to the Shared
    vault.
    """
    env = {k: v for k, v in os.environ.items() if k != "OP_SERVICE_ACCOUNT_TOKEN"}
    result = subprocess.run(
        ["op", "read", OP_ITEM],
        check=True,
        capture_output=True,
        text=True,
        env=env,
    )
    token = result.stdout.strip()
    if not token:
        sys.exit("op read returned an empty token")
    return token


def _http_get_json(url: str, headers: dict[str, str] | None = None) -> Any:
    """Perform a GET request and return the parsed JSON body."""
    req = urllib.request.Request(url, headers=headers or {})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            body = resp.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        sys.exit(f"HTTP {e.code} fetching {url}: {e.read().decode('utf-8', 'replace')}")
    except urllib.error.URLError as e:
        sys.exit(f"Network error fetching {url}: {e.reason}")
    return json.loads(body)


def _list_aliases(token: str) -> dict[str, AliasConfig]:
    """Fetch the alias map from the management API."""
    url = f"{BASE_URL}/v0/management/aliases"
    data = _http_get_json(url, {"x-admin-key": token})
    if not isinstance(data, dict):
        sys.exit(f"Unexpected alias response type: {type(data).__name__}")
    return data


def _lookup_metadata(source: str, source_path: str) -> dict[str, Any]:
    """Fetch normalized metadata for a single alias."""
    params = urllib.parse.urlencode({"source": source, "source_path": source_path})
    url = f"{BASE_URL}/v1/metadata/lookup?{params}"
    payload = _http_get_json(url)
    return payload.get("data") or {}


def _build_model(alias_id: str, meta: dict[str, Any]) -> Model:
    ctx = meta.get("context_length") or 0
    supported = meta.get("supported_parameters") or []
    input_modalities = (meta.get("architecture") or {}).get("input_modalities") or []

    return Model(
        id=alias_id,
        name=clean_name(meta.get("name") or alias_id),
        context_length=ctx,
        max_completion_tokens=(ctx * 20) // 100,
        supports_reasoning="reasoning" in supported,
        supports_image="image" in input_modalities,
    )


def fetch_models() -> list[Model]:
    token = _get_admin_token()
    aliases = _list_aliases(token)

    models: list[Model] = []
    for alias_id, config in sorted(aliases.items()):
        if alias_id in EXCLUDED_IDS:
            continue
        metadata = config.get("metadata")
        if metadata is None:
            continue
        meta = _lookup_metadata(metadata["source"], metadata["source_path"])
        if not meta.get("context_length"):
            continue
        models.append(_build_model(alias_id, meta))

    return models


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    fmt = parser.add_mutually_exclusive_group()
    fmt.add_argument("--tsv", action="store_const", dest="format", const="tsv")
    fmt.add_argument("--json", action="store_const", dest="format", const="json")
    parser.set_defaults(format="tsv")
    args = parser.parse_args()

    models = fetch_models()

    if args.format == "json":
        emit_json(models, sys.stdout)
    else:
        emit_tsv(models, sys.stdout)


if __name__ == "__main__":
    main()