from __future__ import annotations

import json
import subprocess
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

from armorer.openclaw_config import resolve_openclaw_runtime_image_from_compose_dir


def extract_json_object(text: str) -> Optional[Any]:
    raw = str(text or "").strip()
    if not raw:
        return None
    decoder = json.JSONDecoder()
    for idx, ch in enumerate(raw):
        if ch not in "[{":
            continue
        try:
            parsed, _ = decoder.raw_decode(raw[idx:])
            return parsed
        except Exception:
            continue
    return None


def openclaw_runtime_mount_path(install_dir: Path, run_env: Dict[str, str], env_key: str, default_relative: str) -> Path:
    from dotenv import dotenv_values

    raw = str(run_env.get(env_key) or "").strip()
    if not raw:
        env_values = dotenv_values(install_dir / ".env")
        raw = str(env_values.get(env_key) or "").strip()
    if not raw:
        raw = default_relative
    path = Path(raw).expanduser()
    if not path.is_absolute():
        path = (install_dir / path).resolve()
    return path


def openclaw_runtime_image(install_dir: Path, run_env: Dict[str, str]) -> str:
    return resolve_openclaw_runtime_image_from_compose_dir(install_dir, run_env)


def openclaw_provider_models(
    agent: str,
    provider: str,
    *,
    resolve_install_dir: Callable[[str], Optional[Path]],
    build_runtime_env: Callable[[Path], Dict[str, str]],
) -> Tuple[List[Dict[str, Any]], Optional[str]]:
    install_dir = resolve_install_dir(agent)
    if install_dir is None:
        return [], "OpenClaw install directory is unavailable."

    run_env = build_runtime_env(install_dir)
    compose_cmd = [
        "docker",
        "compose",
        "run",
        "--rm",
        "-T",
        "--no-deps",
        "openclaw-cli",
        "models",
        "list",
        "--all",
        "--provider",
        provider,
        "--json",
    ]

    def _run_models(cmd: List[str]) -> subprocess.CompletedProcess[str]:
        return subprocess.run(
            cmd,
            cwd=install_dir,
            env=run_env,
            capture_output=True,
            text=True,
            timeout=120,
        )

    try:
        res = _run_models(compose_cmd)
    except Exception as exc:
        return [], f"OpenClaw model discovery failed: {exc}"

    output = "\n".join(part for part in ((res.stdout or "").strip(), (res.stderr or "").strip()) if part).strip()
    if res.returncode != 0:
        config_dir = openclaw_runtime_mount_path(install_dir, run_env, "OPENCLAW_CONFIG_DIR", ".openclaw")
        workspace_dir = openclaw_runtime_mount_path(install_dir, run_env, "OPENCLAW_WORKSPACE_DIR", "workspace")
        image = openclaw_runtime_image(install_dir, run_env)
        docker_run_cmd = [
            "docker",
            "run",
            "--rm",
            "--network",
            "none",
            "-e",
            "HOME=/home/node",
            "-e",
            "TERM=xterm-256color",
            "-v",
            f"{config_dir}:/home/node/.openclaw",
            "-v",
            f"{workspace_dir}:/home/node/.openclaw/workspace",
            image,
            "node",
            "dist/index.js",
            "models",
            "list",
            "--all",
            "--provider",
            provider,
            "--json",
        ]
        try:
            fallback_res = _run_models(docker_run_cmd)
            fallback_output = "\n".join(
                part for part in ((fallback_res.stdout or "").strip(), (fallback_res.stderr or "").strip()) if part
            ).strip()
            if fallback_res.returncode == 0:
                res = fallback_res
                output = fallback_output
        except Exception:
            pass

    if res.returncode != 0:
        detail = output or f"exit code {res.returncode}"
        return [], f"OpenClaw model discovery failed: {detail[:700]}"
    if "No models found." in output:
        return [], None

    payload = extract_json_object(res.stdout or output)
    if not isinstance(payload, dict):
        return [], "OpenClaw did not return parseable model data."
    models = payload.get("models")
    if not isinstance(models, list):
        return [], "OpenClaw model response did not include a models list."
    return [item for item in models if isinstance(item, dict)], None


def openclaw_discovered_catalog(
    agent: str,
    *,
    resolve_install_dir: Callable[[str], Optional[Path]],
    build_runtime_env: Callable[[Path], Dict[str, str]],
) -> Tuple[List[Dict[str, Any]], Optional[str]]:
    install_dir = resolve_install_dir(agent)
    if install_dir is None:
        return [], "OpenClaw install directory is unavailable."

    run_env = build_runtime_env(install_dir)
    compose_cmd = [
        "docker",
        "compose",
        "run",
        "--rm",
        "-T",
        "--no-deps",
        "openclaw-cli",
        "models",
        "list",
        "--all",
        "--json",
    ]

    def _run_models(cmd: List[str]) -> subprocess.CompletedProcess[str]:
        return subprocess.run(
            cmd,
            cwd=install_dir,
            env=run_env,
            capture_output=True,
            text=True,
            timeout=120,
        )

    try:
        res = _run_models(compose_cmd)
    except Exception as exc:
        return [], f"OpenClaw model catalog discovery failed: {exc}"

    output = "\n".join(part for part in ((res.stdout or "").strip(), (res.stderr or "").strip()) if part).strip()
    if res.returncode != 0:
        config_dir = openclaw_runtime_mount_path(install_dir, run_env, "OPENCLAW_CONFIG_DIR", ".openclaw")
        workspace_dir = openclaw_runtime_mount_path(install_dir, run_env, "OPENCLAW_WORKSPACE_DIR", "workspace")
        image = openclaw_runtime_image(install_dir, run_env)
        docker_run_cmd = [
            "docker",
            "run",
            "--rm",
            "--network",
            "none",
            "-e",
            "HOME=/home/node",
            "-e",
            "TERM=xterm-256color",
            "-v",
            f"{config_dir}:/home/node/.openclaw",
            "-v",
            f"{workspace_dir}:/home/node/.openclaw/workspace",
            image,
            "node",
            "dist/index.js",
            "models",
            "list",
            "--all",
            "--json",
        ]
        try:
            fallback_res = _run_models(docker_run_cmd)
            fallback_output = "\n".join(
                part for part in ((fallback_res.stdout or "").strip(), (fallback_res.stderr or "").strip()) if part
            ).strip()
            if fallback_res.returncode == 0:
                res = fallback_res
                output = fallback_output
        except Exception:
            pass

    if res.returncode != 0:
        detail = output or f"exit code {res.returncode}"
        return [], f"OpenClaw model catalog discovery failed: {detail[:700]}"
    if "No models found." in output:
        return [], None

    payload = extract_json_object(res.stdout or output)
    if not isinstance(payload, dict):
        return [], "OpenClaw did not return parseable model catalog data."
    models = payload.get("models")
    if not isinstance(models, list):
        return [], "OpenClaw model catalog response did not include a models list."
    return [item for item in models if isinstance(item, dict)], None


def openclaw_discovered_providers(models: List[Dict[str, Any]]) -> List[str]:
    providers: List[str] = []
    seen = set()
    for entry in models:
        key = str(entry.get("key") or "").strip()
        if "/" not in key:
            continue
        provider = key.split("/", 1)[0].strip()
        if not provider or provider in seen:
            continue
        seen.add(provider)
        providers.append(provider)
    return providers


def normalize_openclaw_provider_id(provider: str) -> str:
    value = str(provider or "").strip().lower()
    return {
        "google": "gemini",
    }.get(value, value)


def openclaw_model_choice(
    provider: str,
    entry: Dict[str, Any],
    *,
    ui_choice: Callable[[str, Any], Any],
    truncate_menu_text: Callable[[str, int], str],
) -> Optional[Any]:
    key = str(entry.get("key") or "").strip()
    if not key:
        return None
    prefix = f"{provider}/"
    model_id = key[len(prefix):] if key.startswith(prefix) else key
    name = str(entry.get("name") or "").strip()
    hints: List[str] = []
    if name and name != model_id:
        hints.append(name)
    context_window = entry.get("contextWindow")
    if isinstance(context_window, int) and context_window > 0:
        hints.append(f"ctx {context_window}")
    input_mode = str(entry.get("input") or "").strip()
    if input_mode:
        hints.append(input_mode)
    if entry.get("local") is True:
        hints.append("local")
    title = model_id
    if hints:
        title += f"  {' · '.join(hints)}"
    return ui_choice(truncate_menu_text(title, 110), model_id)
