"""Tests for AgentState."""
from dataclasses import dataclass, field
from datetime import UTC, datetime
import json
import sqlite3

import pytest

from logos_adk.agent import Agent
from logos_adk.config import StateConfig
from logos_adk.errors import ConfigError, StateBackendError, StateError
from logos_adk.persistence import get_sqlite_schema_version, sqlite_table_columns
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder
from logos_adk.state import AgentState
from logos_adk.state_store import (
    InMemoryStateStore,
    ShardedStateStore,
    SQLiteStateStore,
    build_state_store,
)
from logos_adk.types import ContentPart, Message, ToolCall, ToolResult


@pytest.fixture(autouse=True)
def _reset_runtime():
    _default_runtime_holder.reset()
    yield
    _default_runtime_holder.reset()


def test_agent_state_defaults():
    state = AgentState()
    assert state.messages == []
    assert state.session_id is None
    assert state.metadata == {}


def test_agent_state_with_values():
    msg = Message(role="user", content="hello")
    state = AgentState(messages=[msg], session_id="s1", metadata={"key": "val"})
    assert len(state.messages) == 1
    assert state.session_id == "s1"
    assert state.metadata["key"] == "val"


def test_agent_state_subclass():
    @dataclass
    class ResearcherState(AgentState):
        last_results: list[dict] = field(default_factory=list)
        sources_visited: set[str] = field(default_factory=set)

    state = ResearcherState()
    assert state.messages == []
    assert state.last_results == []
    assert state.sources_visited == set()
    state.last_results.append({"url": "example.com"})
    assert len(state.last_results) == 1


def test_in_memory_state_store_round_trip_and_rollback():
    store = InMemoryStateStore()
    original = AgentState(
        messages=[Message(role="user", content="hello")],
        session_id="s1",
        metadata={"tenant": "alpha"},
    )

    saved = store.save("bot", original)
    original.messages.append(Message(role="assistant", content="mutated-after-save"))
    loaded = store.load("bot", "s1")

    assert saved.version == 1
    assert loaded is not None
    assert [message.content for message in loaded.messages] == ["hello"]

    store.save(
        "bot",
        AgentState(
            messages=[
                Message(role="user", content="hello"),
                Message(role="assistant", content="world"),
            ],
            session_id="s1",
            metadata={"tenant": "alpha"},
        ),
    )

    history = store.history("bot", "s1")
    exported = store.export_state("bot", "s1", version=1)
    rolled_back = store.rollback("bot", "s1", version=1)

    assert [item["version"] for item in history[:2]] == [2, 1]
    assert exported is not None
    assert exported["version"] == 1
    assert [message.content for message in rolled_back.messages] == ["hello"]
    assert rolled_back.version == 3

    with pytest.raises(StateError, match="No saved state version 99"):
        store.rollback("bot", "s1", version=99)


def test_sharded_state_store_import_uses_payload_session_id_for_routing():
    shard_a = InMemoryStateStore()
    shard_b = InMemoryStateStore()
    store = ShardedStateStore([shard_a, shard_b])

    payload = {
        "state": {
            "messages": [{"role": "user", "content": "restored"}],
            "session_id": "tenant-42",
            "metadata": {"source": "backup"},
            "version": 1,
            "updated_at": None,
        }
    }

    imported = store.import_state("bot", payload)

    assert imported.session_id == "tenant-42"
    assert [message.content for message in imported.messages] == ["restored"]
    assert (shard_a.export_state("bot", "tenant-42") is not None) != (
        shard_b.export_state("bot", "tenant-42") is not None
    )


def test_build_state_store_selects_expected_backends(tmp_path):
    default_store = build_state_store(None)
    persistent_store = build_state_store(StateConfig(), default_persistent=True)
    sqlite_store = build_state_store(
        StateConfig(backend="sqlite", path=str(tmp_path / "state.db"))
    )
    sharded_store = build_state_store(
        StateConfig(
            backend="sharded",
            shards=[
                StateConfig(backend="sqlite", path=str(tmp_path / "a.db")),
                StateConfig(backend="sqlite", path=str(tmp_path / "b.db")),
            ],
        )
    )

    assert isinstance(default_store, InMemoryStateStore)
    assert isinstance(persistent_store, SQLiteStateStore)
    assert isinstance(sqlite_store, SQLiteStateStore)
    assert isinstance(sharded_store, ShardedStateStore)
    assert sharded_store.describe()["shards"] == 2


def test_build_state_store_rejects_invalid_backend_configs(monkeypatch):
    monkeypatch.delenv("DATABASE_URL", raising=False)
    monkeypatch.delenv("LOGOS_ADK_DATABASE_URL", raising=False)

    with pytest.raises(ConfigError, match="requires a non-empty path"):
        build_state_store(StateConfig(backend="sqlite", path=""))

    with pytest.raises(ConfigError, match="requires a DSN"):
        build_state_store(StateConfig(backend="postgresql", dsn=None))

    with pytest.raises(ConfigError, match="requires a non-empty 'shards' list"):
        build_state_store(StateConfig(backend="sharded", shards=[]))

    with pytest.raises(ConfigError, match="Unsupported state backend"):
        build_state_store(StateConfig(backend="mystery"))


def test_sqlite_state_store_connect_wraps_backend_errors(monkeypatch):
    import sqlite3

    def _broken_connect(path: str):
        raise sqlite3.OperationalError("disk i/o error")

    monkeypatch.setattr(sqlite3, "connect", _broken_connect)

    with pytest.raises(StateBackendError, match="Failed to open SQLite state backend"):
        SQLiteStateStore("/tmp/nowhere/state.db").load("bot", "s1")


def test_sqlite_state_store_migrates_legacy_schema(tmp_path):
    db_path = tmp_path / "legacy-state.db"
    state_json = json.dumps(
        {
            "messages": [{"role": "user", "content": "hello"}],
            "session_id": "legacy-session",
            "metadata": {"source": "legacy"},
            "version": 1,
            "updated_at": "2024-01-01T00:00:00+00:00",
        }
    )

    connection = sqlite3.connect(db_path)
    try:
        connection.execute(
            """
            CREATE TABLE agent_state_versions (
                agent_name TEXT NOT NULL,
                session_key TEXT NOT NULL,
                session_id TEXT NOT NULL,
                version INTEGER NOT NULL,
                saved_at TEXT NOT NULL,
                state_json TEXT NOT NULL,
                PRIMARY KEY (agent_name, session_key, version)
            )
            """
        )
        connection.execute(
            """
            INSERT INTO agent_state_versions
            (agent_name, session_key, session_id, version, saved_at, state_json)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            (
                "bot",
                "legacy-session",
                "legacy-session",
                1,
                "2024-01-01T00:00:00+00:00",
                state_json,
            ),
        )
        connection.commit()
    finally:
        connection.close()

    store = SQLiteStateStore(str(db_path))
    loaded = store.load("bot", "legacy-session")
    exported = store.export_state("bot", "legacy-session")
    history = store.history("bot", "legacy-session")

    assert loaded is not None
    assert loaded.session_id == "legacy-session"
    assert [message.content for message in loaded.messages] == ["hello"]
    assert exported is not None
    assert exported["session_id"] == "legacy-session"
    assert exported["state"]["metadata"] == {"source": "legacy"}
    assert history[0]["version"] == 1
    assert history[0]["session_id"] == "legacy-session"

    migrated = sqlite3.connect(db_path)
    try:
        migrated.row_factory = sqlite3.Row
        columns = sqlite_table_columns(migrated, "agent_state_versions")
        assert {"namespace", "session_id", "version", "saved_at", "state_json"} <= columns
        assert "agent_name" not in columns
        assert "session_key" not in columns
        assert get_sqlite_schema_version(migrated, "state_store") == 1
    finally:
        migrated.close()


def test_sqlite_state_store_preserves_tool_payloads_and_content_parts(tmp_path):
    store = SQLiteStateStore(str(tmp_path / "state.db"))
    state = AgentState(
        messages=[
            Message(
                role="assistant",
                content="",
                content_parts=[
                    ContentPart(
                        type="image",
                        uri="https://example.com/chart.png",
                        name="chart",
                        metadata={"width": 640},
                    )
                ],
                tool_calls=[
                    ToolCall(
                        id="call-1",
                        name="search_docs",
                        arguments={"query": "incident summary", "limit": 3},
                    )
                ],
            ),
            Message(
                role="tool",
                content="tool returned",
                tool_result=ToolResult(
                    call_id="call-1",
                    content={"documents": 3},
                    content_parts=[
                        ContentPart(type="text", text="3 documents indexed"),
                        ContentPart(
                            type="file",
                            uri="file:///tmp/report.json",
                            name="report.json",
                        ),
                    ],
                    duration_ms=12.5,
                    status="completed",
                ),
            ),
        ],
        session_id="s-tooling",
        metadata={"checkpointed_at": datetime(2024, 1, 2, tzinfo=UTC)},
    )

    saved = store.save("bot", state)
    loaded = store.load("bot", "s-tooling")
    exported = store.export_state("bot", "s-tooling")

    assert saved.version == 1
    assert loaded is not None
    assert loaded.messages[0].content_parts[0].uri == "https://example.com/chart.png"
    assert loaded.messages[0].content_parts[0].metadata == {"width": 640}
    assert loaded.messages[0].tool_calls[0].arguments == {"query": "incident summary", "limit": 3}
    assert loaded.messages[1].tool_result is not None
    assert loaded.messages[1].tool_result.content == {"documents": 3}
    assert loaded.messages[1].tool_result.content_parts[1].name == "report.json"
    assert exported is not None
    assert exported["state"]["metadata"]["checkpointed_at"] == "2024-01-02T00:00:00+00:00"
    assert exported["state"]["messages"][1]["tool_result"]["content"] == {"documents": 3}


def test_sqlite_state_store_isolates_namespaces_and_clone_imports(tmp_path):
    store = SQLiteStateStore(str(tmp_path / "state.db"))

    store.save(
        "alpha",
        AgentState(
            messages=[Message(role="user", content="alpha-state")],
            session_id="shared-session",
        ),
    )
    store.save(
        "beta",
        AgentState(
            messages=[Message(role="user", content="beta-state")],
            session_id="shared-session",
        ),
    )

    alpha_export = store.export_state("alpha", "shared-session")
    beta_export = store.export_state("beta", "shared-session")
    cloned = store.import_state("alpha", alpha_export or {}, session_id="cloned-session")

    assert alpha_export is not None
    assert beta_export is not None
    assert alpha_export["state"]["messages"][0]["content"] == "alpha-state"
    assert beta_export["state"]["messages"][0]["content"] == "beta-state"
    assert cloned.session_id == "cloned-session"
    assert [message.content for message in store.load("alpha", "cloned-session").messages] == ["alpha-state"]
    assert [message.content for message in store.load("alpha", "shared-session").messages] == ["alpha-state"]
    assert [message.content for message in store.load("beta", "shared-session").messages] == ["beta-state"]
    assert [item["version"] for item in store.history("alpha", "cloned-session")] == [1]


@pytest.mark.asyncio
async def test_state_persists_across_agent_restart(tmp_path):
    store = SQLiteStateStore(str(tmp_path / "state.db"))

    first_provider = MockProvider(responses=["Hello Alex"])
    first_agent = Agent("bot").model(first_provider).state_backend(store)
    await first_agent.run("My name is Alex", session_id="s1")
    _default_runtime_holder.reset()

    second_provider = MockProvider(responses=["Your name is Alex"])
    second_agent = Agent("bot").model(second_provider).state_backend(store)
    await second_agent.run("What is my name?", session_id="s1")

    contents = [message.content for message in second_provider.call_history[0]["messages"]]
    assert "My name is Alex" in contents
    assert "Hello Alex" in contents


@pytest.mark.asyncio
async def test_state_versioning_and_rollback(tmp_path):
    store = SQLiteStateStore(str(tmp_path / "state.db"))
    provider = MockProvider(responses=["first", "second"])
    agent = Agent("bot").model(provider).state_backend(store)

    await agent.run("turn one", session_id="s1")
    await agent.run("turn two", session_id="s1")

    history = agent.list_state_versions("s1")
    assert [item["version"] for item in history[:2]] == [2, 1]

    rolled_back = agent.rollback_state("s1", version=1)
    assert rolled_back.version == 3
    assert [message.content for message in rolled_back.messages] == ["turn one", "first"]


@pytest.mark.asyncio
async def test_state_export_import_round_trip(tmp_path):
    source_store = SQLiteStateStore(str(tmp_path / "source.db"))
    target_store = SQLiteStateStore(str(tmp_path / "target.db"))

    source_agent = Agent("source").model(MockProvider(responses=["saved"])).state_backend(source_store)
    await source_agent.run("remember this", session_id="origin")
    exported = source_agent.export_state("origin")

    assert exported is not None

    target_agent = Agent("target").model(MockProvider(responses=["ok"])).state_backend(target_store)
    imported = target_agent.import_state(exported, session_id="restored")

    assert imported.session_id == "restored"
    assert [message.content for message in imported.messages] == ["remember this", "saved"]


@pytest.mark.asyncio
async def test_agent_state_helpers_expose_backend_info_and_update_hot_sessions():
    store = InMemoryStateStore()
    agent = Agent("bot").model(MockProvider(responses=["saved"])).state_backend(store)

    await agent.run("remember this", session_id="origin")
    exported = agent.export_state("origin")

    assert exported is not None
    assert agent.state_store_info()["backend"] == "in_memory"

    imported = agent.import_state(exported, session_id="clone")

    assert imported.session_id == "clone"
    assert "clone" in agent._sessions
    assert agent.list_state_versions("clone")[0]["version"] == 1


@pytest.mark.asyncio
async def test_sharded_state_store_routes_by_session(tmp_path):
    shard_a = SQLiteStateStore(str(tmp_path / "a.db"))
    shard_b = SQLiteStateStore(str(tmp_path / "b.db"))
    store = ShardedStateStore([shard_a, shard_b])

    agent = Agent("bot").model(MockProvider(responses=["alpha", "beta"])).state_backend(store)
    await agent.run("tenant-a prompt", session_id="tenant-a")
    await agent.run("tenant-b prompt", session_id="tenant-b")

    shard_a_has = shard_a.export_state("bot", "tenant-a") is not None
    shard_b_has = shard_b.export_state("bot", "tenant-a") is not None
    assert shard_a_has != shard_b_has

    shard_a_has = shard_a.export_state("bot", "tenant-b") is not None
    shard_b_has = shard_b.export_state("bot", "tenant-b") is not None
    assert shard_a_has != shard_b_has
