#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 Jiri Vyskocil
# SPDX-License-Identifier: Apache-2.0
# terok:container — this file is deployed into task containers, not used on the host.

"""Unified OpenCode provider launcher for Blablador, KISSKI, and future providers.

This script is self-contained (stdlib-only) and runs inside containers where
terok is NOT installed. Provider-specific configuration is passed via environment
variables injected by the host.

Usage:
  opencode-provider [--list-models] [--base-url BASE_URL] [--]

The script determines the provider name from argv[0] (e.g., "blablador" or "kisski")
and reads TEROK_OC_{NAME}_* environment variables for configuration.
"""

import argparse
import json
import os
import shlex
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any
from urllib import error, request

# Hardcoded fallbacks for manual invocation (when env vars not set)
_FALLBACK_PROVIDERS = {
    "blablador": {
        "display_name": "Helmholtz Blablador",
        "base_url": "https://api.helmholtz-blablador.fz-juelich.de/v1",
        "preferred_model": "alias-huge",
        "fallback_model": "alias-code",
        "env_var_prefix": "BLABLADOR",
        "config_dir": ".blablador",
    },
    "kisski": {
        "display_name": "KISSKI",
        "base_url": "https://chat-ai.academiccloud.de/v1",
        "preferred_model": "devstral-2-123b-instruct-2512",
        "fallback_model": "mistral-large-3-675b-instruct-2512",
        "env_var_prefix": "KISSKI",
        "config_dir": ".kisski",
    },
}


def _resolve_provider_config(name: str | None = None) -> dict[str, str]:
    """Resolve provider configuration from environment variables or fallbacks.

    *name* selects the provider — the basename of ``argv[0]`` (the symlinked
    command) by default, or an explicit ``--provider`` override.  A curated
    provider carries full ``TEROK_OC_{NAME}_*`` config; any other authenticated
    provider the host materialized exposes a generic
    ``TEROK_PROVIDER_{NAME}_BASE_OPENAI_CHAT`` handle, in which case models are
    discovered from the API rather than pinned.
    """
    if name is None:
        name = os.path.basename(sys.argv[0])
    prefix = f"TEROK_OC_{name.upper()}_"

    # Curated provider — full TEROK_OC_* config injected by the host.
    base_url = os.environ.get(f"{prefix}BASE_URL")
    if base_url:
        return {
            "name": name,
            "display_name": os.environ.get(f"{prefix}DISPLAY_NAME", name),
            "base_url": base_url,
            "preferred_model": os.environ.get(f"{prefix}PREFERRED_MODEL", ""),
            "fallback_model": os.environ.get(f"{prefix}FALLBACK_MODEL", ""),
            "env_var_prefix": os.environ.get(f"{prefix}ENV_VAR_PREFIX", name.upper()),
            "config_dir": os.environ.get(f"{prefix}CONFIG_DIR", f".{name}"),
        }

    # Generic provider — any authenticated, openai-chat-compatible provider
    # the host materialized via TEROK_PROVIDER_{NAME}_*.  No pinned models.
    generic_base = os.environ.get(f"TEROK_PROVIDER_{name.upper()}_BASE_OPENAI_CHAT")
    if generic_base:
        return {
            "name": name,
            "display_name": name,
            "base_url": generic_base,
            "preferred_model": "",
            "fallback_model": "",
            "env_var_prefix": name.upper(),
            "config_dir": f".{name}",
        }

    # Fallback to hardcoded defaults for manual invocation
    if name in _FALLBACK_PROVIDERS:
        return {**_FALLBACK_PROVIDERS[name], "name": name}

    raise SystemExit(f"Unknown provider: {name}")


def _config_dir(config: dict[str, str]) -> Path:
    """Return the provider-specific configuration directory."""
    return Path.home() / config["config_dir"]


def _config_path(config: dict[str, str]) -> Path:
    """Return the path to the provider's config.json file."""
    return _config_dir(config) / "config.json"


def _load_api_key(config: dict[str, str]) -> str | None:
    """Load API key from environment or config file.

    Generic providers carry their phantom key in ``TEROK_PROVIDER_{NAME}_TOKEN``;
    curated ones also expose it under ``{PREFIX}_API_KEY``.  Both hold the same
    vault token, so either is fine — the generic var is checked first so an
    arbitrary ``--provider`` selection works without a curated prefix.
    """
    token = os.environ.get(f"TEROK_PROVIDER_{config['name'].upper()}_TOKEN")
    if token:
        return token
    env_var = config["env_var_prefix"] + "_API_KEY"
    api_key = os.environ.get(env_var)
    if api_key:
        return api_key

    cfg_path = _config_path(config)
    if not cfg_path.is_file():
        return None

    try:
        data = json.loads(cfg_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(data, dict):
        return None

    val = data.get("api_key")
    return val if isinstance(val, str) and val.strip() else None


def _fetch_models(base_url: str, api_key: str) -> list[str] | None:
    """Fetch available models from the API. Returns None on failure."""
    url = base_url.rstrip("/") + "/models"
    req = request.Request(
        url,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Accept": "application/json",
        },
    )

    try:
        with request.urlopen(req, timeout=30) as resp:  # nosec B310
            payload = json.loads(resp.read().decode("utf-8"))
    except (error.HTTPError, error.URLError, json.JSONDecodeError):
        return None

    items = []
    if isinstance(payload, dict):
        if isinstance(payload.get("data"), list):
            items = payload["data"]
        elif isinstance(payload.get("models"), list):
            items = payload["models"]

    models: list[str] = []
    for item in items:
        if isinstance(item, dict):
            model_id = item.get("id")
            if isinstance(model_id, str) and model_id:
                models.append(model_id)

    return sorted(set(models)) if models else None


def _build_provider_update(
    config: dict[str, str],
    base_url: str,
    api_key: str,
    model: str,
    models: list[str] | None,
) -> dict[str, Any]:
    """Build the provider-specific config fragment for opencode.json."""
    model_map = {model: {"name": model}}
    if models:
        for mid in models:
            if isinstance(mid, str) and mid:
                model_map.setdefault(mid, {"name": mid})

    provider_name = config["name"]
    return {
        "$schema": "https://opencode.ai/config.json",
        "model": f"{provider_name}/{model}",
        "provider": {
            provider_name: {
                "npm": "@ai-sdk/openai-compatible",
                "name": config["display_name"],
                "options": {
                    "baseURL": base_url,
                    "apiKey": api_key,
                },
                "models": model_map,
            }
        },
        "permission": {
            "*": "allow",
        },
    }


def _merge_provider_config(existing: dict, update: dict, config: dict) -> dict:
    """Merge provider update into existing opencode.json config."""
    merged = dict(existing)

    # Schema handling
    existing_schema = merged.get("$schema")
    expected_schema = update["$schema"]
    if existing_schema and existing_schema != expected_schema:
        print(
            f"Warning: opencode.json has unexpected $schema value "
            f"{existing_schema!r}, expected {expected_schema!r}. Overwriting.",
            file=sys.stderr,
        )
    merged["$schema"] = expected_schema

    # Provider deep-merge
    existing_providers = merged.get("provider")
    if not isinstance(existing_providers, dict):
        existing_providers = {}
    update_providers = update.get("provider", {})
    merged_providers = dict(existing_providers)
    merged_providers.update(update_providers)
    merged["provider"] = merged_providers

    # Model handling - only overwrite if unset or already provider-prefixed
    current_model = merged.get("model")
    if not current_model or (
        isinstance(current_model, str) and current_model.startswith(f"{config['name']}/")
    ):
        merged["model"] = update["model"]

    # Permission - only set if not already configured
    if "permission" not in merged:
        merged["permission"] = update["permission"]

    return merged


def _opencode_config_path(config: dict[str, str]) -> Path:
    """Return the provider-specific OpenCode config path."""
    return _config_dir(config) / "opencode" / "opencode.json"


def _load_opencode_config(config: dict[str, str]) -> dict | None:
    """Load existing OpenCode config if present."""
    config_path = _opencode_config_path(config)
    if not config_path.is_file():
        return None

    try:
        loaded = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    return loaded if isinstance(loaded, dict) else None


def _get_configured_models(config: dict[str, str], existing_config: dict | None) -> set[str]:
    """Extract model IDs from existing provider config."""
    if not existing_config:
        return set()

    try:
        models = existing_config.get("provider", {}).get(config["name"], {}).get("models", {})
        return set(models.keys()) if isinstance(models, dict) else set()
    except (AttributeError, TypeError):
        return set()


def _get_configured_options(config: dict[str, str], existing_config: dict | None) -> dict:
    """Extract provider options from existing config."""
    if not existing_config:
        return {}

    try:
        options = existing_config.get("provider", {}).get(config["name"], {}).get("options", {})
        return options if isinstance(options, dict) else {}
    except (AttributeError, TypeError):
        return {}


def _write_opencode_config(config: dict[str, str], content: dict) -> Path:
    """Write config to OpenCode's location via atomic replace."""
    config_path = _opencode_config_path(config)
    config_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = None

    try:
        with tempfile.NamedTemporaryFile(
            "w", encoding="utf-8", dir=config_path.parent, suffix=".tmp", delete=False
        ) as f:
            tmp_path = f.name
            f.write(json.dumps(content, indent=2) + "\n")
        os.replace(tmp_path, config_path)
        # Security: Restrict config file permissions to owner only
        try:
            os.chmod(config_path, 0o600)
        except (OSError, PermissionError):
            # Permission change may fail in some environments (e.g., read-only filesystems)
            # This is not a security critical failure
            pass
    except BaseException:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)
        raise

    return config_path


def main() -> int:
    """Main entry point for the unified provider launcher."""
    # The provider comes from either the invoked command name (a symlink like
    # ``blablador``) or an explicit ``--provider`` (what the ``opencode`` wrapper
    # passes for a runtime selection).  Resolve it *before* building the full
    # parser, whose help text embeds the resolved provider's config — otherwise
    # ``opencode-provider --provider X`` would resolve the literal command name
    # ``opencode-provider`` and fail with "Unknown provider".
    _pre = argparse.ArgumentParser(add_help=False)
    _pre.add_argument("--provider", default=None)
    _selected, _ = _pre.parse_known_args()
    config = _resolve_provider_config(_selected.provider)

    parser = argparse.ArgumentParser(
        prog=config["name"],
        description=f"Run OpenCode against {config['display_name']} with full permissions.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            f"  {config['name']}                # launch with preferred model\n"
            f"  {config['name']} --list-models  # list available models\n"
            f"  {config['name']} -- --help      # pass flags to opencode\n"
        ),
    )
    parser.add_argument("--list-models", action="store_true", help="List available models and exit")
    parser.add_argument(
        "--base-url",
        default=None,
        help=f"Override API base URL (default: {config['base_url']})",
    )
    parser.add_argument(
        "--provider",
        default=None,
        help="Run against a specific authenticated provider (overrides the command name)",
    )

    args, opencode_args = parser.parse_known_args()
    # ``config`` is already resolved from ``--provider`` (or argv[0]) above; the
    # full parse here just keeps ``--provider`` out of the args forwarded to
    # opencode.

    # Load and validate API key
    api_key = _load_api_key(config)
    if not api_key:
        raise SystemExit(
            f"Missing {config['env_var_prefix']}_API_KEY. Set it in the environment or write "
            f'{{"api_key": "..."}} to {_config_path(config)}.'
        )

    # Determine base URL
    base_url = (
        args.base_url
        or os.environ.get(f"{config['env_var_prefix']}_BASE_URL")
        or config["base_url"]
    )
    base_url = base_url.rstrip("/")

    # Fetch models from API
    fetched_models = _fetch_models(base_url, api_key)

    if args.list_models:
        if fetched_models:
            for model in fetched_models:
                print(model)
        else:
            raise SystemExit(f"Failed to fetch models from {config['display_name']} API")
        return 0

    # Load existing config
    existing_config = _load_opencode_config(config)
    configured_models = _get_configured_models(config, existing_config)

    # Determine which model to use
    model = config["preferred_model"]
    if fetched_models and config["preferred_model"] not in fetched_models:
        fallback = config["fallback_model"]
        if fallback in fetched_models:
            model = fallback
        else:
            # Both defaults gone — pick the first available model
            model = fetched_models[0]
        print(
            f"Warning: Preferred model '{config['preferred_model']}' is no longer available.\n"
            f"Using '{model}'.\n"
            "Check for new upstream versions of terok to update the default model.",
            file=sys.stderr,
        )

    # Update config if needed
    stored_options = _get_configured_options(config, existing_config)
    options_changed = (
        stored_options.get("baseURL") != base_url or stored_options.get("apiKey") != api_key
    )

    if fetched_models:
        fetched_set = set(fetched_models)
        needs_update = fetched_set != configured_models or options_changed
        if needs_update:
            new_models = fetched_set - configured_models
            if new_models:
                print(f"New models available: {', '.join(sorted(new_models))}", file=sys.stderr)
            update = _build_provider_update(config, base_url, api_key, model, fetched_models)
            merged = _merge_provider_config(existing_config or {}, update, config)
            _write_opencode_config(config, merged)
    elif options_changed or not configured_models:
        # Fetch failed: preserve known models if available
        fallback_models = sorted(configured_models) if configured_models else [model]
        update = _build_provider_update(config, base_url, api_key, model, fallback_models)
        merged = _merge_provider_config(existing_config or {}, update, config)
        _write_opencode_config(config, merged)

    # Launch OpenCode
    cmd = ["opencode"] + opencode_args
    env = {**os.environ, "OPENCODE_CONFIG": str(_opencode_config_path(config))}

    # Pinned-alias commands (e.g. ``blablador``) symlink straight to this
    # launcher, bypassing the generated shell wrapper that applies the task's
    # git authorship — so re-apply it here.  The wrapper-driven ``--provider``
    # path (invoked as ``opencode-provider``) has already set the identity.
    if os.path.basename(sys.argv[0]) != "opencode-provider":
        return _launch_with_git_identity(config["name"], cmd, env)

    try:
        return subprocess.call(cmd, env=env)
    except FileNotFoundError:
        raise SystemExit("opencode not found. Rebuild the L1 CLI image to install it.")


_GIT_IDENTITY_HELPER = "/usr/local/share/terok/terok-env-git-identity.sh"
"""Shell helper defining ``_terok_apply_git_identity`` — the single source of the
container's git author/committer policy (staged by the L1 Dockerfile)."""


def _launch_with_git_identity(provider_name: str, cmd: list[str], env: dict[str, str]) -> int:
    """Run *cmd* after applying the task's git authorship for *provider_name*.

    Sources the shared shell helper so a pinned-alias launch commits under the
    same author/committer the generated wrappers would have set — the provider
    name doubles as the agent display name (``blablador`` → ``Blablador``),
    matching the ACP wrapper.  Falls back to a plain launch when the helper is
    absent (standalone images).
    """
    if not os.path.exists(_GIT_IDENTITY_HELPER):
        return subprocess.call(cmd, env=env)
    agent_name = provider_name[:1].upper() + provider_name[1:]
    agent_email = f"noreply@{provider_name}.localhost"
    shim = (
        f". {shlex.quote(_GIT_IDENTITY_HELPER)} && "
        '_terok_apply_git_identity "$1" "$2" && shift 2 && exec "$@"'
    )
    try:
        return subprocess.call(["bash", "-c", shim, "bash", agent_name, agent_email, *cmd], env=env)
    except FileNotFoundError:
        raise SystemExit("bash not found; cannot apply git identity for the pinned alias.")


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