from pydantic import BaseModel

from haiku.rag.tools.context import ToolContext


class TestState(BaseModel):
    value: int = 0


class TestStateWithList(BaseModel):
    items: list[str] = []


def test_tool_context_defaults():
    """Test ToolContext has sensible defaults."""
    ctx = ToolContext()
    assert ctx._namespaces == {}


def test_register_and_get():
    """Test register and get state for a namespace."""
    ctx = ToolContext()
    state = TestState(value=42)

    ctx.register("test.namespace", state)

    retrieved = ctx.get("test.namespace")
    assert retrieved is state
    assert retrieved.value == 42


def test_get_nonexistent_namespace():
    """Test get returns None for unregistered namespace."""
    ctx = ToolContext()
    assert ctx.get("nonexistent") is None


def test_get_or_create_creates_new():
    """Test get_or_create creates state when namespace doesn't exist."""
    ctx = ToolContext()

    state = ctx.get_or_create("test.namespace", TestStateWithList)
    assert isinstance(state, TestStateWithList)
    assert state.items == []


def test_get_or_create_returns_existing():
    """Test get_or_create returns existing state."""
    ctx = ToolContext()

    state1 = ctx.get_or_create("test.namespace", TestStateWithList)
    state1.items.append("item1")

    state2 = ctx.get_or_create("test.namespace", TestStateWithList)
    assert state2 is state1
    assert state2.items == ["item1"]


def test_clear_namespace():
    """Test clear_namespace removes only the specified namespace."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=1))
    ctx.register("ns2", TestState(value=2))

    ctx.clear_namespace("ns1")

    assert ctx.get("ns1") is None
    ns2 = ctx.get("ns2")
    assert isinstance(ns2, TestState)
    assert ns2.value == 2


def test_clear_namespace_nonexistent():
    """Test clear_namespace handles nonexistent namespace gracefully."""
    ctx = ToolContext()
    ctx.clear_namespace("nonexistent")  # Should not raise


def test_clear_all():
    """Test clear_all clears all namespaces."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=1))
    ctx.register("ns2", TestState(value=2))

    ctx.clear_all()

    assert ctx.get("ns1") is None
    assert ctx.get("ns2") is None


def test_namespaces_property():
    """Test namespaces property lists all registered namespaces."""
    ctx = ToolContext()
    assert ctx.namespaces == []

    ctx.register("ns1", TestState())
    ctx.register("ns2", TestState())

    assert set(ctx.namespaces) == {"ns1", "ns2"}


def test_shared_namespace_between_toolsets():
    """Test that toolsets can share state via the same namespace."""

    class SharedState(BaseModel):
        citations: dict[str, int] = {}

    SHARED_NAMESPACE = "haiku.rag.citations"

    ctx = ToolContext()

    # First toolset registers the shared state
    state1 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
    state1.citations["chunk-a"] = 1

    # Second toolset gets the same state
    state2 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
    assert state2 is state1
    assert state2.citations == {"chunk-a": 1}

    # Both see updates
    state2.citations["chunk-b"] = 2
    assert state1.citations == {"chunk-a": 1, "chunk-b": 2}


def test_dump_namespaces():
    """Test dump_namespaces serializes all registered states."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=10))
    ctx.register("ns2", TestStateWithList(items=["a", "b"]))

    data = ctx.dump_namespaces()

    assert data == {
        "ns1": {"value": 10},
        "ns2": {"items": ["a", "b"]},
    }


def test_load_namespace():
    """Test load_namespace deserializes and registers state."""
    ctx = ToolContext()

    state = ctx.load_namespace("ns1", TestState, {"value": 42})

    assert isinstance(state, TestState)
    assert state.value == 42
    assert ctx.get("ns1") is state


def test_serialization_roundtrip():
    """Test full serialization/deserialization roundtrip."""
    # Create and populate context
    original = ToolContext()
    original.register("search", TestStateWithList(items=["result1", "result2"]))
    original.register("qa", TestState(value=99))

    # Serialize
    ns_data = original.dump_namespaces()

    # Deserialize
    restored = ToolContext()
    restored.load_namespace("search", TestStateWithList, ns_data["search"])
    restored.load_namespace("qa", TestState, ns_data["qa"])

    # Verify
    search_state = restored.get("search")
    assert isinstance(search_state, TestStateWithList)
    assert search_state.items == ["result1", "result2"]

    qa_state = restored.get("qa")
    assert isinstance(qa_state, TestState)
    assert qa_state.value == 99


def test_get_with_type_match():
    """Test get with state_type returns typed state when type matches."""
    ctx = ToolContext()
    state = TestState(value=42)
    ctx.register("ns", state)

    result = ctx.get("ns", TestState)
    assert result is state
    assert result.value == 42


def test_get_with_type_mismatch():
    """Test get with state_type returns None when type doesn't match."""
    ctx = ToolContext()
    ctx.register("ns", TestState(value=42))

    result = ctx.get("ns", TestStateWithList)
    assert result is None


def test_get_without_type():
    """Test get without state_type returns BaseModel (unchanged behavior)."""
    ctx = ToolContext()
    state = TestState(value=42)
    ctx.register("ns", state)

    result = ctx.get("ns")
    assert result is state


def test_tool_context_state_key_default_none():
    """Test ToolContext state_key defaults to None."""
    ctx = ToolContext()
    assert ctx.state_key is None


def test_tool_context_state_key_set():
    """Test ToolContext state_key can be set."""
    ctx = ToolContext()
    ctx.state_key = "haiku.rag.chat"
    assert ctx.state_key == "haiku.rag.chat"


# =============================================================================
# ToolContextCache Tests
# =============================================================================


def test_tool_context_cache_get_or_create_new():
    """Test get_or_create returns a new context with is_new=True."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    context, is_new = cache.get_or_create("thread-1")

    assert isinstance(context, ToolContext)
    assert is_new is True


def test_tool_context_cache_get_or_create_existing():
    """Test get_or_create returns existing context with is_new=False."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    ctx1, is_new1 = cache.get_or_create("thread-1")
    ctx2, is_new2 = cache.get_or_create("thread-1")

    assert ctx2 is ctx1
    assert is_new1 is True
    assert is_new2 is False


def test_tool_context_cache_different_keys():
    """Test get_or_create returns different contexts for different keys."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    ctx1, _ = cache.get_or_create("thread-1")
    ctx2, _ = cache.get_or_create("thread-2")

    assert ctx1 is not ctx2


def test_tool_context_cache_ttl_expiry():
    """Test that contexts are evicted after TTL expires."""
    from datetime import timedelta

    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache(ttl=timedelta(seconds=0))
    ctx1, _ = cache.get_or_create("thread-1")

    # With zero TTL, next access should create a new context
    ctx2, is_new = cache.get_or_create("thread-1")

    assert ctx2 is not ctx1
    assert is_new is True


def test_tool_context_cache_remove():
    """Test remove deletes a specific key."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    cache.get_or_create("thread-1")
    cache.get_or_create("thread-2")

    cache.remove("thread-1")

    ctx, is_new = cache.get_or_create("thread-1")
    assert is_new is True

    # thread-2 should still exist
    ctx2, is_new2 = cache.get_or_create("thread-2")
    assert is_new2 is False


def test_tool_context_cache_remove_nonexistent():
    """Test remove handles nonexistent key gracefully."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    cache.remove("nonexistent")  # Should not raise


def test_tool_context_cache_clear():
    """Test clear removes all entries."""
    from haiku.rag.tools.context import ToolContextCache

    cache = ToolContextCache()
    cache.get_or_create("thread-1")
    cache.get_or_create("thread-2")

    cache.clear()

    ctx1, is_new1 = cache.get_or_create("thread-1")
    ctx2, is_new2 = cache.get_or_create("thread-2")
    assert is_new1 is True
    assert is_new2 is True


# =============================================================================
# build_state_snapshot / restore_state_snapshot Tests
# =============================================================================


class NestedModel(BaseModel):
    name: str = ""
    count: int = 0


class StateWithNested(BaseModel):
    nested: NestedModel | None = None
    tags: list[str] = []


def test_build_state_snapshot_empty():
    """build_state_snapshot on empty context returns empty dict."""
    ctx = ToolContext()
    assert ctx.build_state_snapshot() == {}


def test_build_state_snapshot_single_namespace():
    """build_state_snapshot with one namespace returns its fields."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=42))
    snapshot = ctx.build_state_snapshot()
    assert snapshot == {"value": 42}


def test_build_state_snapshot_multiple_namespaces():
    """build_state_snapshot merges fields from all namespaces."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=42))
    ctx.register("ns2", TestStateWithList(items=["a", "b"]))
    snapshot = ctx.build_state_snapshot()
    assert snapshot == {"value": 42, "items": ["a", "b"]}


def test_build_state_snapshot_nested_model():
    """build_state_snapshot serializes nested models with mode='json'."""
    from datetime import datetime

    class TimestampState(BaseModel):
        ts: datetime | None = None

    ctx = ToolContext()
    ctx.register("ns", TimestampState(ts=datetime(2025, 1, 27, 12, 0, 0)))
    snapshot = ctx.build_state_snapshot()
    assert isinstance(snapshot["ts"], str)
    assert snapshot["ts"] == "2025-01-27T12:00:00"


def test_restore_state_snapshot_empty_context():
    """restore_state_snapshot on empty context is a no-op."""
    ctx = ToolContext()
    ctx.restore_state_snapshot({"value": 42})
    assert ctx.namespaces == []


def test_restore_state_snapshot_single_namespace():
    """restore_state_snapshot updates matching fields in registered namespaces."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=0))
    ctx.restore_state_snapshot({"value": 99})
    state = ctx.get("ns1", TestState)
    assert state is not None
    assert state.value == 99


def test_restore_state_snapshot_partial_update():
    """restore_state_snapshot only touches fields present in data."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=42))
    ctx.register("ns2", TestStateWithList(items=["original"]))
    # Only update ns2's items, not ns1's value
    ctx.restore_state_snapshot({"items": ["updated"]})
    ns1 = ctx.get("ns1", TestState)
    ns2 = ctx.get("ns2", TestStateWithList)
    assert ns1 is not None
    assert ns2 is not None
    assert ns1.value == 42
    assert ns2.items == ["updated"]


def test_restore_state_snapshot_nested_model():
    """restore_state_snapshot deserializes nested models from dicts."""
    ctx = ToolContext()
    ctx.register("ns", StateWithNested())
    ctx.restore_state_snapshot({"nested": {"name": "foo", "count": 5}, "tags": ["x"]})
    state = ctx.get("ns", StateWithNested)
    assert state is not None
    assert state.nested is not None
    assert state.nested.name == "foo"
    assert state.nested.count == 5
    assert state.tags == ["x"]


def test_state_snapshot_roundtrip():
    """build then restore produces equivalent state."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=42))
    ctx.register("ns2", TestStateWithList(items=["a", "b"]))

    snapshot = ctx.build_state_snapshot()

    ctx2 = ToolContext()
    ctx2.register("ns1", TestState())
    ctx2.register("ns2", TestStateWithList())
    ctx2.restore_state_snapshot(snapshot)

    ns1 = ctx2.get("ns1", TestState)
    ns2 = ctx2.get("ns2", TestStateWithList)
    assert ns1 is not None
    assert ns2 is not None
    assert ns1.value == 42
    assert ns2.items == ["a", "b"]


def test_restore_state_snapshot_ignores_unknown_fields():
    """restore_state_snapshot ignores fields not in any registered namespace."""
    ctx = ToolContext()
    ctx.register("ns1", TestState(value=0))
    ctx.restore_state_snapshot({"value": 10, "unknown_field": "ignored"})
    ns1 = ctx.get("ns1", TestState)
    assert ns1 is not None
    assert ns1.value == 10