"""JSON schema validation evaluation rule.

Validates that the agent reply is valid JSON and optionally conforms to
a specified structural schema (key presence, types, nested paths).

Configuration params:
    required_keys   list of dot-path keys that must exist (e.g. ["data.id", "status"])
    forbidden_keys  list of dot-path keys that must NOT exist
    key_types       dict mapping dot-path key → expected type name
                    ("str", "int", "float", "bool", "list", "dict", "null")
                    e.g. {"status": "str", "count": "int"}
    strict_json     if true (default), fail if reply is not parseable JSON

All params are optional; with no params the rule only checks valid JSON.
"""

import json
from typing import Any

from agenteval.evaluation.rules.base import EvalRule, RuleResult, register_rule
from agenteval.models import Case, Turn
from agenteval.utils.llm import extract_reply_text

_TYPE_MAP: dict[str, type] = {
    "str": str,
    "int": int,
    "float": float,
    "bool": bool,
    "list": list,
    "dict": dict,
    "null": type(None),
}


def _get_path(data: Any, path: str) -> tuple[bool, Any]:
    """Return (found, value) for a dot-separated path."""
    parts = path.split(".")
    current = data
    for part in parts:
        if isinstance(current, dict):
            if part not in current:
                return False, None
            current = current[part]
        elif isinstance(current, list) and part.isdigit():
            idx = int(part)
            if idx >= len(current):
                return False, None
            current = current[idx]
        else:
            return False, None
    return True, current


@register_rule
class JsonSchemaRule(EvalRule):
    """Validate that the reply is valid JSON and matches a structural schema."""

    name = "json_schema"

    async def evaluate(self, case: Case, dialog: list[Turn]) -> RuleResult:
        if not dialog:
            return RuleResult(passed=False, reason="无回复记录")

        reply_text = extract_reply_text(dialog[-1].reply).strip()
        strict_json: bool = self.params.get("strict_json", True)
        required_keys: list[str] = self.params.get("required_keys", [])
        forbidden_keys: list[str] = self.params.get("forbidden_keys", [])
        key_types: dict[str, str] = self.params.get("key_types", {})

        # Try to find JSON in the reply (may be wrapped in markdown code block)
        data: Any = None
        try:
            # Strip markdown code fences if present
            text = reply_text
            if text.startswith("```"):
                lines = text.split("\n")
                text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
            data = json.loads(text)
        except (json.JSONDecodeError, ValueError):
            if strict_json:
                return RuleResult(passed=False, reason="回复不是合法 JSON")
            # Non-strict: proceed with None data, required_keys will catch it

        errors: list[str] = []

        if data is not None:
            # Check required keys
            for key in required_keys:
                found, _ = _get_path(data, key)
                if not found:
                    errors.append(f"缺少字段: {key}")

            # Check forbidden keys
            for key in forbidden_keys:
                found, _ = _get_path(data, key)
                if found:
                    errors.append(f"存在禁止字段: {key}")

            # Check key types
            for key, expected_type_name in key_types.items():
                found, value = _get_path(data, key)
                if not found:
                    errors.append(f"类型检查字段缺失: {key}")
                    continue
                expected_type = _TYPE_MAP.get(expected_type_name)
                if expected_type is None:
                    continue  # unknown type name — skip
                if not isinstance(value, expected_type):
                    actual = type(value).__name__
                    errors.append(f"{key} 类型错误: 期望 {expected_type_name}，实际 {actual}")

        if errors:
            return RuleResult(passed=False, reason="; ".join(errors))

        total_checks = len(required_keys) + len(forbidden_keys) + len(key_types)
        return RuleResult(
            passed=True,
            score=1.0,
            reason=f"JSON 结构校验通过（{total_checks} 项检查）" if total_checks else "合法 JSON",
        )