# Copyright (c) Microsoft. All rights reserved.

"""Tests for type definitions in _types.py."""

import pytest
from pydantic import ValidationError

from agent_framework_ag_ui._types import AgentState, AGUIRequest, PredictStateConfig, RunMetadata


class TestPredictStateConfig:
    """Test PredictStateConfig TypedDict."""

    def test_predict_state_config_creation(self) -> None:
        """Test creating a PredictStateConfig dict."""
        config: PredictStateConfig = {
            "state_key": "document",
            "tool": "write_document",
            "tool_argument": "content",
        }

        assert config["state_key"] == "document"
        assert config["tool"] == "write_document"
        assert config["tool_argument"] == "content"

    def test_predict_state_config_with_none_tool_argument(self) -> None:
        """Test PredictStateConfig with None tool_argument."""
        config: PredictStateConfig = {
            "state_key": "status",
            "tool": "update_status",
            "tool_argument": None,
        }

        assert config["state_key"] == "status"
        assert config["tool"] == "update_status"
        assert config["tool_argument"] is None

    def test_predict_state_config_type_validation(self) -> None:
        """Test that PredictStateConfig validates field types at runtime."""
        config: PredictStateConfig = {
            "state_key": "test",
            "tool": "test_tool",
            "tool_argument": "arg",
        }

        assert isinstance(config["state_key"], str)
        assert isinstance(config["tool"], str)
        assert isinstance(config["tool_argument"], (str, type(None)))


class TestRunMetadata:
    """Test RunMetadata TypedDict."""

    def test_run_metadata_creation(self) -> None:
        """Test creating a RunMetadata dict."""
        metadata: RunMetadata = {
            "run_id": "run-123",
            "thread_id": "thread-456",
            "predict_state": [
                {
                    "state_key": "document",
                    "tool": "write_document",
                    "tool_argument": "content",
                }
            ],
        }

        assert metadata["run_id"] == "run-123"
        assert metadata["thread_id"] == "thread-456"
        assert metadata["predict_state"] is not None
        assert len(metadata["predict_state"]) == 1
        assert metadata["predict_state"][0]["state_key"] == "document"

    def test_run_metadata_with_none_predict_state(self) -> None:
        """Test RunMetadata with None predict_state."""
        metadata: RunMetadata = {
            "run_id": "run-789",
            "thread_id": "thread-012",
            "predict_state": None,
        }

        assert metadata["run_id"] == "run-789"
        assert metadata["thread_id"] == "thread-012"
        assert metadata["predict_state"] is None

    def test_run_metadata_empty_predict_state(self) -> None:
        """Test RunMetadata with empty predict_state list."""
        metadata: RunMetadata = {
            "run_id": "run-345",
            "thread_id": "thread-678",
            "predict_state": [],
        }

        assert metadata["run_id"] == "run-345"
        assert metadata["thread_id"] == "thread-678"
        assert metadata["predict_state"] == []


class TestAgentState:
    """Test AgentState TypedDict."""

    def test_agent_state_creation(self) -> None:
        """Test creating an AgentState dict."""
        state: AgentState = {
            "messages": [
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there"},
            ]
        }

        assert state["messages"] is not None
        assert len(state["messages"]) == 2
        assert state["messages"][0]["role"] == "user"
        assert state["messages"][1]["role"] == "assistant"

    def test_agent_state_with_none_messages(self) -> None:
        """Test AgentState with None messages."""
        state: AgentState = {"messages": None}

        assert state["messages"] is None

    def test_agent_state_empty_messages(self) -> None:
        """Test AgentState with empty messages list."""
        state: AgentState = {"messages": []}

        assert state["messages"] == []

    def test_agent_state_complex_messages(self) -> None:
        """Test AgentState with complex message structures."""
        state: AgentState = {
            "messages": [
                {
                    "role": "user",
                    "content": "Test",
                    "metadata": {"timestamp": "2025-10-30"},
                },
                {
                    "role": "assistant",
                    "content": "Response",
                    "tool_calls": [{"name": "search", "args": {}}],
                },
            ]
        }

        assert state["messages"] is not None
        assert len(state["messages"]) == 2
        assert "metadata" in state["messages"][0]
        assert "tool_calls" in state["messages"][1]


class TestAGUIRequest:
    """Test AGUIRequest Pydantic model."""

    def test_agui_request_minimal(self) -> None:
        """Test creating AGUIRequest with only required fields."""
        request = AGUIRequest.model_validate({"messages": [{"role": "user", "content": "Hello"}]})

        assert len(request.messages) == 1
        assert request.messages[0]["content"] == "Hello"
        assert request.run_id is None
        assert request.thread_id is None
        assert request.state is None
        assert request.tools is None
        assert request.context is None
        assert request.forwarded_props is None
        assert request.parent_run_id is None

    def test_agui_request_all_fields(self) -> None:
        """Test creating AGUIRequest with all fields populated."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "run_id": "run-123",
                "thread_id": "thread-456",
                "state": {"counter": 0},
                "tools": [{"name": "search", "description": "Search tool"}],
                "context": [{"type": "document", "content": "Some context"}],
                "forwarded_props": {"custom_key": "custom_value"},
                "parent_run_id": "parent-run-789",
            }
        )

        assert request.run_id == "run-123"
        assert request.thread_id == "thread-456"
        assert request.state == {"counter": 0}
        assert request.tools == [{"name": "search", "description": "Search tool"}]
        assert request.context == [{"type": "document", "content": "Some context"}]
        assert request.forwarded_props == {"custom_key": "custom_value"}
        assert request.parent_run_id == "parent-run-789"

    def test_agui_request_camel_case_aliases(self) -> None:
        """Test AGUIRequest accepts camelCase aliases from AG-UI HTTP clients."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "runId": "run-camel-1",
                "threadId": "thread-camel-1",
                "forwardedProps": {"k": "v"},
                "parentRunId": "parent-camel-1",
            }
        )

        assert request.run_id == "run-camel-1"
        assert request.thread_id == "thread-camel-1"
        assert request.forwarded_props == {"k": "v"}
        assert request.parent_run_id == "parent-camel-1"

    def test_agui_request_model_dump_excludes_none(self) -> None:
        """Test that model_dump(exclude_none=True) excludes None fields."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "test"}],
                "tools": [{"name": "my_tool"}],
                "context": [{"id": "ctx1"}],
            }
        )

        dumped = request.model_dump(exclude_none=True)

        assert "messages" in dumped
        assert "tools" in dumped
        assert "context" in dumped
        assert "run_id" not in dumped
        assert "thread_id" not in dumped
        assert "state" not in dumped
        assert "forwarded_props" not in dumped
        assert "parent_run_id" not in dumped

    def test_agui_request_model_dump_includes_all_set_fields(self) -> None:
        """Test that model_dump preserves all explicitly set fields.

        This is critical for the fix - ensuring tools, context, forwarded_props,
        and parent_run_id are not stripped during request validation.
        """
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "test"}],
                "tools": [{"name": "client_tool", "parameters": {"type": "object"}}],
                "context": [{"type": "snippet", "content": "code here"}],
                "forwarded_props": {"auth_token": "secret", "user_id": "user-1"},
                "parent_run_id": "parent-456",
            }
        )

        dumped = request.model_dump(exclude_none=True)

        # Verify all fields are preserved (the main bug fix)
        assert dumped["tools"] == [{"name": "client_tool", "parameters": {"type": "object"}}]
        assert dumped["context"] == [{"type": "snippet", "content": "code here"}]
        assert dumped["forwarded_props"] == {"auth_token": "secret", "user_id": "user-1"}
        assert dumped["parent_run_id"] == "parent-456"

    def test_agui_request_available_interrupts_alias_round_trip(self) -> None:
        """availableInterrupts should deserialize to canonical Interrupt models."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "availableInterrupts": [{"id": "req_1", "reason": "input_required", "message": "Choose"}],
            }
        )

        assert request.available_interrupts is not None
        assert request.available_interrupts[0].id == "req_1"
        assert request.available_interrupts[0].reason == "input_required"
        dumped = request.model_dump(exclude_none=True)
        assert dumped["available_interrupts"] == [{"id": "req_1", "reason": "input_required", "message": "Choose"}]
        assert "availableInterrupts" not in dumped

    def test_agui_request_resume_accepts_canonical_entries(self) -> None:
        """resume should preserve AG-UI resume arrays at the HTTP trust boundary."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "resume": [{"interruptId": "req_1", "status": "resolved", "payload": {"approved": True}}],
            }
        )

        assert request.resume is not None
        assert request.resume[0].interrupt_id == "req_1"
        assert request.resume[0].status == "resolved"
        assert request.resume[0].payload == {"approved": True}

    def test_agui_request_resume_schema_advertises_canonical_entries(self) -> None:
        """resume should advertise the canonical ResumeEntry array shape in JSON schema."""
        resume_schema = AGUIRequest.model_json_schema()["properties"]["resume"]
        array_schema = next((schema for schema in resume_schema["anyOf"] if schema.get("type") == "array"), None)

        assert array_schema is not None
        assert array_schema["items"] == {"$ref": "#/$defs/ResumeEntry"}

    def test_agui_request_resume_accepts_legacy_object_shapes(self) -> None:
        """resume coerces supported legacy containers to canonical ResumeEntry models."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "resume": {"interrupts": [{"id": "req_1", "value": {"approved": True}}]},
            }
        )

        assert request.resume is not None
        assert request.resume[0].interrupt_id == "req_1"
        assert request.resume[0].status == "resolved"
        assert request.resume[0].payload == {"approved": True}

    def test_agui_request_resume_accepts_legacy_single_entry_mapping(self) -> None:
        """resume coerces a supported single legacy entry object to a one-entry canonical list."""
        request = AGUIRequest.model_validate(
            {
                "messages": [{"role": "user", "content": "Hello"}],
                "resume": {"toolCallId": "call_1", "approved": True},
            }
        )

        assert request.resume is not None
        assert request.resume[0].interrupt_id == "call_1"
        assert request.resume[0].status == "resolved"
        assert request.resume[0].payload == {"approved": True}

    def test_agui_request_resume_rejects_malformed_shape(self) -> None:
        """resume rejects malformed inputs at request validation once the contract shape is advertised."""
        with pytest.raises(ValidationError):
            AGUIRequest.model_validate(
                {
                    "messages": [{"role": "user", "content": "Hello"}],
                    "resume": {"unexpected": "shape"},
                }
            )
