# Copyright (C) 2026 GoodData Corporation

"""Run an external command that returns JSON to stdout.

This is a building block for fully decoupling validation tool implementations from AIDA.
An "external tool" can live in another repository and be invoked via a CommandSpec from
`.aida/validation_registry.yaml` (or its included files).
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, cast

import orjson

from aida_mcp.tools.core.validation_registry import CommandSpec, get_validation_registry
from aida_mcp.tools.validate.stream_progress import ProgressContext, build_forwarder
from aida_mcp.utils.command_runner import CommandResult, run_command_streaming

logger = logging.getLogger(__name__)


def _format_token(token: str, values: dict[str, str]) -> str:
    """Best-effort placeholder expansion (unknown placeholders remain intact)."""
    try:
        return token.format_map(values)
    except KeyError:
        return token
    except Exception:
        return token


def _format_argv(argv: list[str], values: dict[str, str]) -> list[str]:
    return [_format_token(tok, values) for tok in argv]


def _json_loads(raw: str) -> object:
    return orjson.loads(raw)


def _parse_json_from_stdout(stdout: str) -> dict[str, Any] | None:
    text = (stdout or "").strip()
    if not text:
        return None
    # First try full stdout.
    try:
        obj = _json_loads(text)
        if isinstance(obj, dict):
            return cast(dict[str, Any], obj)
    except Exception:
        pass
    # Fallback: try last non-empty line (allows tools to print logs before JSON).
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if not lines:
        return None
    try:
        obj = _json_loads(lines[-1])
    except Exception:
        return None
    if not isinstance(obj, dict):
        return None
    return cast(dict[str, Any], obj)


async def run_external_json_tool(
    *,
    tool_name: str,
    command_id: str,
    values: dict[str, str],
    workspace_root: Path,
    timeout_sec: int | None = None,
    extra_env: dict[str, str] | None = None,
    ctx: ProgressContext | None = None,
) -> tuple[dict[str, Any], CommandResult]:
    """Run an external tool (from registry) and parse its JSON output.

    Returns:
        (payload, cmd_result) where payload is always a dict (success or failure).
    """
    registry = get_validation_registry()
    spec = registry.commands.get(command_id)
    if not isinstance(spec, CommandSpec):
        return (
            {
                "tool": tool_name,
                "success": False,
                "text": f"External runner command_id not found in registry: {command_id}",
                "error_excerpt": f"command_id not found: {command_id}",
                "error_count": 1,
                "log_paths": [],
            },
            CommandResult(exit_code=127, stdout="", stderr="command_id not found", log_path=None),
        )

    cmd = _format_argv(spec.argv, values)
    cwd_str = _format_token(spec.cwd, values) if isinstance(spec.cwd, str) and spec.cwd else str(workspace_root)
    cwd = Path(cwd_str)

    env = dict(spec.env)
    if extra_env:
        env.update(extra_env)

    effective_timeout = timeout_sec if timeout_sec is not None else spec.timeout_sec
    timeout = int(effective_timeout) if isinstance(effective_timeout, int) and effective_timeout > 0 else 3600

    # Prefer streaming so we can forward progress events if present.
    forwarder = build_forwarder(
        ctx=ctx,
        current_step=1,
        total_steps=1,
        step_label=f"{tool_name} (external)",
        enable_tail_fallback=False,
    )
    cmd_result = await run_command_streaming(
        cmd,
        cwd=cwd,
        timeout=timeout,
        env=env or None,
        on_stdout_line=forwarder.on_stdout_line,
        on_stderr_line=forwarder.on_stderr_line,
        ctx=None,  # avoid duplicate ctx.info from command runner here
        step_name=None,
    )

    parsed = _parse_json_from_stdout(cmd_result.stdout)
    if isinstance(parsed, dict):
        parsed.setdefault("tool", tool_name)
        return parsed, cmd_result

    # Invalid JSON -> return a structured error but keep logs.
    excerpt = (cmd_result.stderr or cmd_result.stdout or "").strip()
    if len(excerpt) > 800:
        excerpt = excerpt[:800] + "..."
    payload: dict[str, Any] = {
        "tool": tool_name,
        "success": False,
        "text": f"External {tool_name} runner did not return JSON output (exit code: {cmd_result.exit_code})",
        "error_excerpt": excerpt or "no output",
        "error_count": 1,
        "log_paths": [str(cmd_result.log_path)] if cmd_result.log_path else [],
        "command": " ".join(cmd),
        "exit_code": cmd_result.exit_code,
    }
    return payload, cmd_result
