#!/usr/bin/env python3
"""Render registry JSON files as markdown tables for local developer reference.

Usage:
    python render.py                  # print all three registries
    python render.py embedders
    python render.py llms
    python render.py rerankers
    python render.py embedders llms
"""

import json
import sys
from pathlib import Path

REGISTRIES_DIR = Path(__file__).parent / "goodmem" / "registries"

PROVIDER_DISPLAY = {
    "OPENAI": "OpenAI",
    "COHERE": "Cohere",
    "JINA": "Jina",
    "VOYAGE": "Voyage AI",
    "ANTHROPIC": "Anthropic",
    "GOOGLE": "Google",
    "MISTRAL": "Mistral",
}

PROVIDER_ENDPOINTS = {
    "OPENAI": "https://api.openai.com/v1",
    "COHERE": "https://api.cohere.com/v2",
    "JINA": "https://api.jina.ai/v1",
    "VOYAGE": "https://api.voyageai.com/v1",
    "ANTHROPIC": "https://api.anthropic.com/v1",
    "GOOGLE": "https://generativelanguage.googleapis.com/v1",
    "MISTRAL": "https://api.mistral.ai/v1",
}


def load(name: str) -> dict[str, dict]:
    with open(REGISTRIES_DIR / f"{name}.json") as f:
        return json.load(f)


def fmt_num(n: int | None) -> str:
    return f"{n:,}" if n is not None else "--"


def fmt_modalities(modalities: list[str]) -> str:
    return ", ".join(m.lower() for m in modalities)


def providers_in_order(registry: dict) -> list[str]:
    seen: list[str] = []
    for entry in registry.values():
        p = entry.get("provider_type", "")
        if p and p not in seen:
            seen.append(p)
    return seen


def render_embedders(registry: dict) -> str:
    lines = ["# Embedders\n"]
    for provider_type in providers_in_order(registry):
        display = PROVIDER_DISPLAY.get(provider_type, provider_type)
        endpoint = PROVIDER_ENDPOINTS.get(provider_type)
        lines.append(f"\n## {display}\n")
        if endpoint:
            lines.append(f"**API Endpoint:** `{endpoint}`\n")
        lines.append("| Model | Modalities | Languages | Default Dim | Dimensions | Max Tokens | Special Features | Status |")
        lines.append("|-------|------------|-----------|-------------|------------|------------|------------------|--------|")
        for model_id, e in registry.items():
            if e.get("provider_type") != provider_type:
                continue
            link = e.get("link", "")
            name_col = f"[{model_id}]({link})" if link else model_id
            dims = e.get("dimensions", {})
            default_dim = dims.get("default", "--") if dims else "--"
            if "allowed" in dims:
                dim_info = ", ".join(str(d) for d in dims["allowed"])
            elif "range" in dims:
                dim_info = f"{dims['range'][0]}–{dims['range'][1]}"
            else:
                dim_info = "--"
            lines.append(
                f"| {name_col}"
                f" | {fmt_modalities(e.get('supported_modalities', []))}"
                f" | {e.get('languages', '--')}"
                f" | {default_dim}"
                f" | {dim_info}"
                f" | {fmt_num(e.get('max_sequence_length'))}"
                f" | {e.get('special_features', '--')}"
                f" | {e.get('status', '--')} |"
            )
    return "\n".join(lines) + "\n"


def render_llms(registry: dict) -> str:
    lines = ["# LLMs\n"]
    for provider_type in providers_in_order(registry):
        display = PROVIDER_DISPLAY.get(provider_type, provider_type)
        endpoint = PROVIDER_ENDPOINTS.get(provider_type)
        lines.append(f"\n## {display}\n")
        if endpoint:
            lines.append(f"**API Endpoint:** `{endpoint}`\n")
        lines.append("| Model | Modalities | Context Window | Special Features | Status |")
        lines.append("|-------|------------|----------------|------------------|--------|")
        for model_id, e in registry.items():
            if e.get("provider_type") != provider_type:
                continue
            link = e.get("link", "")
            name_col = f"[{model_id}]({link})" if link else model_id
            lines.append(
                f"| {name_col}"
                f" | {fmt_modalities(e.get('supported_modalities', []))}"
                f" | {fmt_num(e.get('max_context_length'))}"
                f" | {e.get('special_features', '--')}"
                f" | {e.get('status', '--')} |"
            )
    return "\n".join(lines) + "\n"


def render_rerankers(registry: dict) -> str:
    lines = ["# Rerankers\n"]
    for provider_type in providers_in_order(registry):
        display = PROVIDER_DISPLAY.get(provider_type, provider_type)
        endpoint = PROVIDER_ENDPOINTS.get(provider_type)
        lines.append(f"\n## {display}\n")
        if endpoint:
            lines.append(f"**API Endpoint:** `{endpoint}`\n")
        lines.append("| Model | Languages | Max Tokens | Special Features | Status |")
        lines.append("|-------|-----------|------------|------------------|--------|")
        for model_id, e in registry.items():
            if e.get("provider_type") != provider_type:
                continue
            link = e.get("link", "")
            name_col = f"[{model_id}]({link})" if link else model_id
            lines.append(
                f"| {name_col}"
                f" | {e.get('languages', '--')}"
                f" | {fmt_num(e.get('max_sequence_length'))}"
                f" | {e.get('special_features', '--')}"
                f" | {e.get('status', '--')} |"
            )
    return "\n".join(lines) + "\n"


RENDERERS = {
    "embedders": render_embedders,
    "llms": render_llms,
    "rerankers": render_rerankers,
}

ALL = ["embedders", "llms", "rerankers"]

names = sys.argv[1:] or ALL
unknown = [n for n in names if n not in RENDERERS]
if unknown:
    print(f"Unknown registry: {', '.join(unknown)}. Choose from: {', '.join(ALL)}", file=sys.stderr)
    sys.exit(1)

for i, name in enumerate(names):
    if i > 0:
        print("\n---\n")
    print(RENDERERS[name](load(name)))
