#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Validate analyzer JSON output against schemas/analyze-data.schema.json.

This intentionally implements only the JSON Schema keywords used by the local
schema so it can run without third-party dependencies in benchmark environments.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, List, Optional, Sequence


class ValidationError(Exception):
    """Raised for schema loading or schema reference errors."""


def default_schema_path() -> str:
    script_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.abspath(os.path.join(script_dir, "..", "schemas", "analyze-data.schema.json"))


def load_json(path: str) -> Any:
    try:
        with open(path, "r", encoding="utf-8") as handle:
            return json.load(handle)
    except FileNotFoundError as exc:
        raise ValidationError(f"file not found: {path}") from exc
    except json.JSONDecodeError as exc:
        raise ValidationError(f"invalid JSON in {path}: line {exc.lineno}, column {exc.colno}: {exc.msg}") from exc
    except OSError as exc:
        raise ValidationError(f"unable to read {path}: {exc}") from exc


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Validate TensorRT analysis JSON generated by analyze_trt_perf.py."
    )
    parser.add_argument("input", help="Path to generated analysis JSON.")
    parser.add_argument(
        "--schema",
        default=default_schema_path(),
        help="Schema path. Defaults to schemas/analyze-data.schema.json.",
    )
    parser.add_argument(
        "--max-errors",
        type=int,
        default=40,
        help="Maximum validation errors to print. Defaults to 40.",
    )
    return parser.parse_args(argv)


def decode_ref_token(token: str) -> str:
    return token.replace("~1", "/").replace("~0", "~")


def resolve_ref(ref: str, root_schema: Any) -> Any:
    if not ref.startswith("#/"):
        raise ValidationError(f"only local JSON Schema refs are supported: {ref}")

    current = root_schema
    for token in ref[2:].split("/"):
        key = decode_ref_token(token)
        if not isinstance(current, dict) or key not in current:
            raise ValidationError(f"unresolved schema ref: {ref}")
        current = current[key]
    return current


def type_label(instance: Any) -> str:
    if instance is None:
        return "null"
    if isinstance(instance, bool):
        return "boolean"
    if isinstance(instance, int):
        return "integer"
    if isinstance(instance, float):
        return "number"
    if isinstance(instance, str):
        return "string"
    if isinstance(instance, list):
        return "array"
    if isinstance(instance, dict):
        return "object"
    return type(instance).__name__


def matches_type(instance: Any, expected: str) -> bool:
    if expected == "null":
        return instance is None
    if expected == "boolean":
        return isinstance(instance, bool)
    if expected == "integer":
        return isinstance(instance, int) and not isinstance(instance, bool)
    if expected == "number":
        return isinstance(instance, (int, float)) and not isinstance(instance, bool)
    if expected == "string":
        return isinstance(instance, str)
    if expected == "array":
        return isinstance(instance, list)
    if expected == "object":
        return isinstance(instance, dict)
    raise ValidationError(f"unsupported schema type: {expected}")


def matches_any_type(instance: Any, expected: Any) -> bool:
    if isinstance(expected, str):
        return matches_type(instance, expected)
    if isinstance(expected, list):
        return any(matches_type(instance, item) for item in expected)
    raise ValidationError(f"invalid schema type declaration: {expected!r}")


def append_error(errors: List[str], path: str, message: str) -> None:
    errors.append(f"{path}: {message}")


def append_nearest_errors(errors: List[str], matches: List[List[str]], limit: int = 5) -> None:
    if not matches:
        return
    nearest = min(matches, key=lambda item: (sum("expected constant" in error for error in item), len(item)))
    for error in nearest[:limit]:
        errors.append(f"  nearest schema: {error}")


def validate(instance: Any, schema: Any, path: str, root_schema: Any) -> List[str]:
    if not isinstance(schema, dict):
        raise ValidationError(f"{path}: schema node is not an object")

    if "$ref" in schema:
        return validate(instance, resolve_ref(str(schema["$ref"]), root_schema), path, root_schema)

    errors: List[str] = []

    if "allOf" in schema:
        for subschema in schema["allOf"]:
            errors.extend(validate(instance, subschema, path, root_schema))

    if "anyOf" in schema:
        matches = [validate(instance, subschema, path, root_schema) for subschema in schema["anyOf"]]
        if not any(not match_errors for match_errors in matches):
            append_error(errors, path, "does not match any allowed schema")
            append_nearest_errors(errors, matches)

    if "oneOf" in schema:
        matches = [validate(instance, subschema, path, root_schema) for subschema in schema["oneOf"]]
        match_count = sum(1 for match_errors in matches if not match_errors)
        if match_count != 1:
            append_error(errors, path, f"matches {match_count} schemas, expected exactly one")
            if match_count == 0:
                append_nearest_errors(errors, matches)

    if "const" in schema and instance != schema["const"]:
        append_error(errors, path, f"expected constant {schema['const']!r}, got {instance!r}")

    if "enum" in schema and instance not in schema["enum"]:
        append_error(errors, path, f"expected one of {schema['enum']!r}, got {instance!r}")

    expected_type = schema.get("type")
    if expected_type is not None and not matches_any_type(instance, expected_type):
        append_error(errors, path, f"expected type {expected_type!r}, got {type_label(instance)}")
        return errors

    if isinstance(instance, dict):
        validate_object(instance, schema, path, root_schema, errors)
    elif isinstance(instance, list):
        validate_array(instance, schema, path, root_schema, errors)
    elif isinstance(instance, (int, float)) and not isinstance(instance, bool):
        validate_number(instance, schema, path, errors)

    return errors


def validate_object(instance: dict, schema: dict, path: str, root_schema: Any, errors: List[str]) -> None:
    properties = schema.get("properties", {})
    required = schema.get("required", [])
    additional = schema.get("additionalProperties", True)

    for key in required:
        if key not in instance:
            append_error(errors, path, f"missing required property {key!r}")

    for key, value in instance.items():
        child_path = f"{path}.{key}"
        if key in properties:
            errors.extend(validate(value, properties[key], child_path, root_schema))
        elif additional is False:
            append_error(errors, child_path, "additional property is not allowed")
        elif isinstance(additional, dict):
            errors.extend(validate(value, additional, child_path, root_schema))


def validate_array(instance: list, schema: dict, path: str, root_schema: Any, errors: List[str]) -> None:
    min_items = schema.get("minItems")
    if isinstance(min_items, int) and len(instance) < min_items:
        append_error(errors, path, f"expected at least {min_items} items, got {len(instance)}")

    max_items = schema.get("maxItems")
    if isinstance(max_items, int) and len(instance) > max_items:
        append_error(errors, path, f"expected at most {max_items} items, got {len(instance)}")

    items_schema = schema.get("items")
    if isinstance(items_schema, dict):
        for index, value in enumerate(instance):
            errors.extend(validate(value, items_schema, f"{path}[{index}]", root_schema))


def validate_number(instance: float, schema: dict, path: str, errors: List[str]) -> None:
    if "minimum" in schema and instance < schema["minimum"]:
        append_error(errors, path, f"expected >= {schema['minimum']}, got {instance}")
    if "maximum" in schema and instance > schema["maximum"]:
        append_error(errors, path, f"expected <= {schema['maximum']}, got {instance}")


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)

    try:
        data = load_json(args.input)
        schema = load_json(args.schema)
        errors = validate(data, schema, "$", schema)
    except ValidationError as exc:
        sys.stderr.write(f"error: {exc}\n")
        return 2

    if errors:
        sys.stderr.write(f"{args.input}: FAIL ({len(errors)} schema error(s))\n")
        for error in errors[: args.max_errors]:
            sys.stderr.write(f"- {error}\n")
        remaining = len(errors) - args.max_errors
        if remaining > 0:
            sys.stderr.write(f"... {remaining} more error(s) omitted\n")
        return 1

    sys.stdout.write(f"{args.input}: PASS\n")
    return 0


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