# Copyright 2026 Datastrato, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Verify that the Python Pydantic protocol covers and is consistent with the ADP protocol schema.

Schema is the single source of truth (schema/adp-protocol-{version}.json). These tests ensure:
- Coverage: every schema $defs type (except excluded) has a Python type
- Enum consistency: schema enum values match Python Literal type alias args
- Required fields and property existence: object types have matching required/optional fields
"""

# pyright: reportCallIssue=false

import json
import typing
import unittest
from pathlib import Path
from typing import Any

from jsonschema import Draft202012Validator
from pydantic import BaseModel

from adp_sdk import types as types_module
from adp_sdk.types import (
    LATEST_PROTOCOL_VERSION,
    Capabilities,
    DescribeResult,
    ExecuteResult,
    Implementation,
    InitializeResult,
    ServerCapabilities,
    UsageContract,
)

# Schema path: schema/adp-protocol-{version}.json at project root
_SCHEMA_DIR = Path(__file__).resolve().parent.parent / "schema"
SCHEMA_PATH = _SCHEMA_DIR / f"adp-protocol-{LATEST_PROTOCOL_VERSION}.json"

# Schema $defs keys that do not require a Python type
# - Generic / alias-only types that map to Python builtins or have no dedicated class
# - Full JSON-RPC request wrappers that the SDK doesn't have as dedicated types
SCHEMA_DEFS_EXCLUDED = {
    "Request",
    "Record<string,unknown>",
    "PingRequest",
    "DiscoverRequest",
    "DescribeRequest",
    "ExecuteRequest",
    "ValidateRequest",
    "InitializeRequest",
    # Type aliases that map to Python builtins (str, str | int)
    "Cursor",
    "ResourceId",
    "TraceId",
    "ProgressToken",
    "RequestId",
}

# Map schema type name -> Python type name (when different)
SCHEMA_TO_PYTHON_NAME: dict[str, str] = {
    "EmptyResult": "Result",
    "Field": "FieldDefinition",
}


class TestProtocolSchemaConsistency(unittest.TestCase):
    """
    Verify protocol types and the versioned ADP protocol schema are in sync.

    Public methods (test_*) are the test cases. Private methods (_*) are
    helpers used only within this class.
    """

    _schema: dict[str, Any]

    @classmethod
    def setUpClass(cls) -> None:
        cls._schema = cls._load_schema()

    # -------------------------------------------------------------------------
    # Public: test methods
    # -------------------------------------------------------------------------

    def test_schema_file_exists(self) -> None:
        """Schema file exists at expected path."""
        self.assertTrue(SCHEMA_PATH.exists(), f"Schema not found at {SCHEMA_PATH}")

    def test_protocol_coverage_every_schema_def_has_python_type(self) -> None:
        """
        For each schema $defs key (except excluded), the SDK exposes a corresponding type.
        """
        schema = self._schema
        missing: list[str] = []
        for name in self._schema_def_names(schema, exclude=SCHEMA_DEFS_EXCLUDED):
            py_type = self._get_python_type(name)
            if py_type is None:
                missing.append(name)
        self.assertFalse(
            missing,
            (
                f"Schema types with no Python equivalent: {missing}. "
                "Add corresponding types to adp_sdk.types or add to SCHEMA_DEFS_EXCLUDED."
            ),
        )

    def test_enum_consistency_schema_enum_matches_python_literal(self) -> None:
        """
        For each schema definition that is a string enum, the Python Literal
        type alias has the same set of values.
        """
        schema = self._schema
        defs = self._get_schema_defs(schema)
        for name, defn in defs.items():
            if not self._is_enum_def(defn):
                continue
            schema_values: set[str] = set(defn["enum"])
            py_type = self._get_python_type(name)
            self.assertIsNotNone(py_type, f"Schema enum {name} has no Python type (coverage should have failed)")
            # SDK uses Literal type aliases for enums
            origin = typing.get_origin(py_type)
            if origin is typing.Literal:
                py_values: set[str] = set(typing.get_args(py_type))
                self.assertEqual(
                    schema_values,
                    py_values,
                    (
                        f"Enum {name}: schema has {schema_values}, Python has {py_values}. "
                        f"Symdiff: {schema_values ^ py_values}"
                    ),
                )

    def test_object_type_required_fields_and_properties_match_schema(self) -> None:
        """
        For each schema object type with properties, the Python model has
        the same required fields (by protocol name) and at least the same properties.
        """
        schema = self._schema
        defs = self._get_schema_defs(schema)
        for name, defn in defs.items():
            if not self._is_object_def(defn):
                continue
            py_type = self._get_python_type(name)
            if py_type is None:
                continue
            if not isinstance(py_type, type) or not issubclass(py_type, BaseModel):
                continue

            schema_props: set[str] = set(defn.get("properties", {}))
            schema_required: set[str] = set(defn.get("required", []))
            py_protocol_names = self._get_protocol_field_names(py_type)

            missing_props = schema_props - py_protocol_names
            self.assertFalse(
                missing_props,
                (
                    f"Object type {name}: schema properties {missing_props} missing in Python "
                    f"model. Protocol field names: {py_protocol_names}"
                ),
            )

            missing_required_as_props = schema_required - py_protocol_names
            self.assertFalse(
                missing_required_as_props,
                (
                    f"Object type {name}: schema required fields "
                    f"{missing_required_as_props} missing in Python model. "
                    f"Protocol field names: {py_protocol_names}"
                ),
            )

    def test_roundtrip_initialize_result_valid_against_schema(self) -> None:
        """Serialized InitializeResult (minimal valid instance) validates against schema."""
        schema = self._schema
        obj = InitializeResult(
            protocol_version="2026-01-20",
            capabilities=ServerCapabilities(),
            server_info=Implementation(name="test", version="0.1.0"),
        )
        payload = obj.model_dump(by_alias=True, exclude_none=True)

        ref_schema = self._schema_for_def(schema, "InitializeResult")
        Draft202012Validator(ref_schema).validate(payload)  # type: ignore[no-untyped-call]

    def test_roundtrip_describe_result_valid_against_schema(self) -> None:
        """Serialized DescribeResult (minimal valid instance) validates against schema."""
        schema = self._schema
        obj = DescribeResult.model_validate(
            {
                "resourceId": "domain:alias",
                "version": 1,
                "intentClass": "QUERY",
                "usageContract": UsageContract(capabilities=Capabilities(), fields=[]).model_dump(by_alias=True),
            }
        )
        payload = obj.model_dump(by_alias=True, exclude_none=True)

        ref_schema = self._schema_for_def(schema, "DescribeResult")
        Draft202012Validator(ref_schema).validate(payload)  # type: ignore[no-untyped-call]

    def test_roundtrip_execute_result_valid_against_schema(self) -> None:
        """Serialized ExecuteResult (minimal valid instance) validates against schema."""
        schema = self._schema
        obj = ExecuteResult(results=[])
        payload = obj.model_dump(by_alias=True, exclude_none=True)

        ref_schema = self._schema_for_def(schema, "ExecuteResult")
        Draft202012Validator(ref_schema).validate(payload)  # type: ignore[no-untyped-call]

    # -------------------------------------------------------------------------
    # Private: helpers (internal to this test class)
    # -------------------------------------------------------------------------

    @staticmethod
    def _load_schema() -> dict[str, Any]:
        """Load schema.json from project root."""
        with open(SCHEMA_PATH, encoding="utf-8") as f:
            return json.load(f)  # type: ignore[no-any-return]

    @staticmethod
    def _get_schema_defs(schema: dict[str, Any]) -> dict[str, Any]:
        """Return $defs from schema."""
        return schema.get("$defs", {})  # type: ignore[no-any-return]

    def _schema_def_names(self, schema: dict[str, Any], *, exclude: set[str] | None = None) -> list[str]:
        """Return sorted list of $defs keys, excluding given names."""
        defs = self._get_schema_defs(schema)
        excl = exclude or set()
        return sorted(k for k in defs if k not in excl)

    @staticmethod
    def _get_python_type(schema_name: str) -> Any:
        """
        Resolve schema type name to Python type from adp_sdk.types.
        Uses SCHEMA_TO_PYTHON_NAME for aliases (e.g. EmptyResult -> Result).
        """
        python_name = SCHEMA_TO_PYTHON_NAME.get(schema_name, schema_name)
        return getattr(types_module, python_name, None)

    @staticmethod
    def _get_protocol_field_names(model: type[BaseModel]) -> set[str]:
        """
        Return the set of protocol (wire) field names for a Pydantic model.
        Uses explicit alias if set, otherwise falls back to the Python field name.
        """
        names: set[str] = set()
        for name, info in model.model_fields.items():
            alias = info.alias
            if alias is not None:
                names.add(alias)
            else:
                names.add(name)
        return names

    @staticmethod
    def _is_object_def(defn: dict[str, Any]) -> bool:
        """True if schema definition is an object type with properties."""
        return defn.get("type") == "object" and "properties" in defn

    @staticmethod
    def _is_enum_def(defn: dict[str, Any]) -> bool:
        """True if schema definition is a string enum."""
        return defn.get("type") == "string" and "enum" in defn

    @staticmethod
    def _schema_for_def(schema: dict[str, Any], type_name: str) -> dict[str, Any]:
        """Build a schema that validates against $defs[type_name] with $ref resolution."""
        return {
            "$defs": schema["$defs"],
            "$ref": f"#/$defs/{type_name}",
        }
