# Copyright (C) 2026 GoodData Corporation

"""Guard for tool responses to keep MCP payloads small."""

from __future__ import annotations

from datetime import UTC, datetime
from hashlib import sha256
from typing import Any, cast

import orjson

from aida_mcp.output.log_utils import truncate_output_line
from aida_mcp.output.policy import MAX_MCP_RESPONSE_BYTES, MAX_RESULT_LINE_CHARS
from aida_mcp.utils.artifacts import get_artifacts_root


def shrink_tool_payload(payload: Any) -> dict[str, Any]:
    """Ensure response payload stays under MAX_RESULT_BYTES."""
    if not isinstance(payload, dict):
        final = {"success": False, "text": str(payload)[:200]}
        _write_payload_snapshot(final)
        return final

    payload_dict = cast(dict[str, Any], payload)

    def _payload_size(candidate: dict[str, Any]) -> int | None:
        try:
            return len(orjson.dumps(candidate, option=orjson.OPT_NON_STR_KEYS))
        except (TypeError, ValueError):
            return None

    size = _payload_size(payload_dict)
    if size is None:
        final = {"tool": payload_dict.get("tool"), "success": False, "text": "Invalid response payload"}
        _write_payload_snapshot(final)
        return final

    if size <= MAX_MCP_RESPONSE_BYTES:
        text_value = payload_dict.get("text")
        if isinstance(text_value, str) and text_value:
            safe_lines = [truncate_output_line(line, MAX_RESULT_LINE_CHARS) for line in text_value.splitlines()]
            payload_dict["text"] = "\n".join(safe_lines)
        _write_payload_snapshot(payload_dict)
        return payload_dict

    log_paths = payload_dict.get("log_paths")
    errors = payload_dict.get("errors")
    error_count = payload_dict.get("error_count")
    steps = payload_dict.get("steps")

    step_dicts: list[dict[str, Any]] = []
    if isinstance(steps, list):
        steps_list = cast(list[object], steps)
        for step in steps_list:
            if isinstance(step, dict):
                step_dicts.append(cast(dict[str, Any], step))

    if not log_paths and step_dicts:
        step_logs: list[str] = []
        for step in step_dicts:
            log_path = step.get("log_path")
            if isinstance(log_path, str) and log_path:
                step_logs.append(log_path)
        if step_logs:
            log_paths = step_logs

    if errors is None and step_dicts:
        step_errors: list[str] = []
        for step in step_dicts:
            excerpt = step.get("error_excerpt")
            if isinstance(excerpt, str) and excerpt:
                step_errors.append(excerpt)
        if step_errors:
            errors = step_errors

    error_excerpt = None
    if isinstance(errors, list) and errors:
        errors_list_raw = cast(list[object], errors)
        errors_list = [str(line) for line in errors_list_raw if line is not None]
        excerpt_lines = [line for line in errors_list[-3:] if line]
        if excerpt_lines:
            error_excerpt = "\n".join(excerpt_lines)

    minimal: dict[str, Any] = {
        "tool": payload_dict.get("tool"),
        "success": payload_dict.get("success"),
        "text": "Response too large; see log_paths.",
    }
    if log_paths:
        minimal["log_paths"] = log_paths
    if error_excerpt:
        minimal["error_excerpt"] = truncate_output_line(error_excerpt, MAX_RESULT_LINE_CHARS)
    if error_count is not None:
        minimal["error_count"] = error_count

    minimal_size = _payload_size(minimal) or 0
    if minimal_size > MAX_MCP_RESPONSE_BYTES:
        if "log_paths" in minimal and isinstance(minimal["log_paths"], list):
            minimal["log_paths"] = minimal["log_paths"][:3]
        if "error_excerpt" in minimal:
            minimal["error_excerpt"] = truncate_output_line(str(minimal["error_excerpt"]), 200)
        minimal_size = _payload_size(minimal) or 0
        if minimal_size > MAX_MCP_RESPONSE_BYTES:
            minimal.pop("error_excerpt", None)
            minimal.pop("error_count", None)
            minimal_size = _payload_size(minimal) or 0
            if minimal_size > MAX_MCP_RESPONSE_BYTES:
                minimal = {
                    "tool": payload_dict.get("tool"),
                    "success": payload_dict.get("success"),
                    "text": "Response too large; see log_paths.",
                }
    _write_payload_snapshot(minimal)
    return minimal


def _write_payload_snapshot(payload: dict[str, Any]) -> None:
    """Write the exact JSON payload that will be sent to the MCP client."""
    try:
        content = orjson.dumps(payload, option=orjson.OPT_NON_STR_KEYS).decode("utf-8")
    except (TypeError, ValueError):
        return
    try:
        root = get_artifacts_root() / "mcp-responses"
        root.mkdir(parents=True, exist_ok=True)
        timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S-%f")
        digest = sha256(content.encode("utf-8")).hexdigest()[:12]
        path = root / f"{timestamp}-{digest}.json"
        path.write_text(content, encoding="utf-8")
    except OSError:
        return
