from __future__ import annotations

import uuid
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock, Mock

import pytest

from tracecat.agent.session.schemas import AgentSessionCreate, AgentSessionUpdate
from tracecat.agent.session.service import AgentSessionService
from tracecat.agent.session.types import AgentSessionEntity
from tracecat.agent.subagents import ResolvedAgentsConfig
from tracecat.auth.types import Role
from tracecat.db.models import AgentSession
from tracecat.exceptions import TracecatValidationError


def _build_service() -> tuple[AgentSessionService, SimpleNamespace, Role]:
    workspace_id = uuid.uuid4()
    role = Role(
        type="service",
        service_id="tracecat-api",
        workspace_id=workspace_id,
        organization_id=uuid.uuid4(),
        scopes=frozenset({"agent:execute"}),
    )
    session = SimpleNamespace(
        add=Mock(),
        commit=AsyncMock(),
        refresh=AsyncMock(),
    )
    return AgentSessionService(cast(Any, session), role), session, role


@pytest.mark.anyio
async def test_create_session_preserves_null_preset_version_for_current() -> None:
    service, session, _role = _build_service()
    preset_id = uuid.uuid4()
    agent_session_id = uuid.uuid4()
    validate_mock = AsyncMock(return_value=None)
    agents_binding_mock = AsyncMock(return_value=None)
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    created = await service.create_session(
        AgentSessionCreate(
            id=agent_session_id,
            title="Chat",
            entity_type=AgentSessionEntity.AGENT_PRESET,
            entity_id=preset_id,
            agent_preset_version_id=None,
        )
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.AGENT_PRESET,
        entity_id=preset_id,
        agent_preset_id=None,
        agent_preset_version_id=None,
    )
    agents_binding_mock.assert_awaited_once_with(None)
    assert created.agent_preset_id == preset_id
    assert created.agent_preset_version_id is None
    assert created.agents_binding is None
    session.add.assert_called_once_with(created)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(created)


@pytest.mark.anyio
async def test_create_session_derives_agents_binding_from_pinned_preset_version() -> (
    None
):
    service, session, _role = _build_service()
    preset_id = uuid.uuid4()
    pinned_version_id = uuid.uuid4()
    agents_binding = {"enabled": True, "subagents": []}
    validate_mock = AsyncMock(return_value=pinned_version_id)
    agents_binding_mock = AsyncMock(return_value=agents_binding)
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    created = await service.create_session(
        AgentSessionCreate(
            title="Chat",
            entity_type=AgentSessionEntity.CASE,
            entity_id=uuid.uuid4(),
            agent_preset_id=preset_id,
            agent_preset_version_id=pinned_version_id,
        )
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.CASE,
        entity_id=created.entity_id,
        agent_preset_id=preset_id,
        agent_preset_version_id=pinned_version_id,
    )
    assert created.agent_preset_id == preset_id
    assert created.agent_preset_version_id == pinned_version_id
    assert created.agents_binding == agents_binding
    agents_binding_mock.assert_awaited_once_with(pinned_version_id)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(created)


@pytest.mark.anyio
async def test_create_session_prefers_provided_agents_binding_for_pinned_preset() -> (
    None
):
    service, session, _role = _build_service()
    preset_id = uuid.uuid4()
    pinned_version_id = uuid.uuid4()
    child_preset_id = uuid.uuid4()
    child_version_id = uuid.uuid4()
    validate_mock = AsyncMock(return_value=pinned_version_id)
    agents_binding_mock = AsyncMock(
        side_effect=AssertionError("should not derive binding when provided")
    )
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock
    agents_binding = ResolvedAgentsConfig.model_validate(
        {
            "enabled": True,
            "subagents": [
                {
                    "preset": "child",
                    "preset_version": 1,
                    "preset_id": child_preset_id,
                    "preset_version_id": child_version_id,
                }
            ],
        }
    )

    created = await service.create_session(
        AgentSessionCreate(
            title="Chat",
            entity_type=AgentSessionEntity.CASE,
            entity_id=uuid.uuid4(),
            agent_preset_id=preset_id,
            agent_preset_version_id=pinned_version_id,
        ),
        agents_binding=agents_binding,
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.CASE,
        entity_id=created.entity_id,
        agent_preset_id=preset_id,
        agent_preset_version_id=pinned_version_id,
    )
    assert created.agent_preset_id == preset_id
    assert created.agent_preset_version_id == pinned_version_id
    assert created.agents_binding == agents_binding.model_dump(mode="json")
    agents_binding_mock.assert_not_called()
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(created)


@pytest.mark.anyio
async def test_create_session_persists_internal_agents_binding_without_preset() -> None:
    service, session, _role = _build_service()
    validate_mock = AsyncMock(return_value=None)
    agents_binding_mock = AsyncMock(return_value=None)
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock
    agents_binding = ResolvedAgentsConfig.model_validate(
        {"enabled": True, "subagents": []}
    )

    created = await service.create_session(
        AgentSessionCreate(
            title="Chat",
            entity_type=AgentSessionEntity.CASE,
            entity_id=uuid.uuid4(),
        ),
        agents_binding=agents_binding,
    )

    assert created.agents_binding == {"enabled": True, "subagents": []}
    agents_binding_mock.assert_not_called()
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(created)


@pytest.mark.anyio
async def test_update_session_preserves_null_version_when_preset_changes() -> None:
    service, session, role = _build_service()
    old_preset_id = uuid.uuid4()
    new_preset_id = uuid.uuid4()
    old_version_id = uuid.uuid4()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="case",
        entity_id=uuid.uuid4(),
        agent_preset_id=old_preset_id,
        agent_preset_version_id=old_version_id,
        agents_binding={"enabled": True, "subagents": []},
    )
    validate_mock = AsyncMock(return_value=None)
    agents_binding_mock = AsyncMock(return_value=None)
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(agent_preset_id=new_preset_id),
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.CASE,
        entity_id=agent_session.entity_id,
        agent_preset_id=new_preset_id,
        agent_preset_version_id=None,
    )
    assert updated.agent_preset_id == new_preset_id
    assert updated.agent_preset_version_id is None
    assert updated.agents_binding is None
    agents_binding_mock.assert_awaited_once_with(None)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_update_session_clears_agents_binding_when_preset_removed() -> None:
    service, session, role = _build_service()
    old_preset_id = uuid.uuid4()
    old_version_id = uuid.uuid4()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="case",
        entity_id=uuid.uuid4(),
        agent_preset_id=old_preset_id,
        agent_preset_version_id=old_version_id,
        agents_binding={"enabled": True, "subagents": []},
    )

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(agent_preset_id=None),
    )

    assert updated.agent_preset_id is None
    assert updated.agent_preset_version_id is None
    assert updated.agents_binding is None
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_validate_preset_assignment_preserves_null_without_resolving_current() -> (
    None
):
    service, _session, _role = _build_service()
    preset_id = uuid.uuid4()

    preset_service = Mock()
    preset_service.get_preset = AsyncMock(return_value=SimpleNamespace(id=preset_id))
    preset_service.resolve_agent_preset_version = AsyncMock()

    with pytest.MonkeyPatch.context() as mp:
        mp.setattr(
            "tracecat.agent.session.service.AgentPresetService",
            Mock(return_value=preset_service),
        )

        pinned_version_id = await service._validate_preset_version_for_assignment(
            entity_type=AgentSessionEntity.AGENT_PRESET,
            entity_id=preset_id,
            agent_preset_id=None,
            agent_preset_version_id=None,
        )

    assert pinned_version_id is None
    preset_service.get_preset.assert_awaited_once_with(preset_id)
    preset_service.resolve_agent_preset_version.assert_not_awaited()


@pytest.mark.anyio
async def test_update_session_allows_version_only_repin_for_preset_sessions() -> None:
    service, session, role = _build_service()
    preset_id = uuid.uuid4()
    new_version_id = uuid.uuid4()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="agent_preset",
        entity_id=preset_id,
        agent_preset_id=None,
        agent_preset_version_id=uuid.uuid4(),
    )
    validate_mock = AsyncMock(return_value=new_version_id)
    agents_binding_mock = AsyncMock(return_value={"enabled": True, "subagents": []})
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(agent_preset_version_id=new_version_id),
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.AGENT_PRESET,
        entity_id=preset_id,
        agent_preset_id=None,
        agent_preset_version_id=new_version_id,
    )
    assert updated.agent_preset_id == preset_id
    assert updated.agent_preset_version_id == new_version_id
    assert updated.agents_binding == {"enabled": True, "subagents": []}
    agents_binding_mock.assert_awaited_once_with(new_version_id)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_update_session_clears_pinned_version_to_follow_current() -> None:
    service, session, role = _build_service()
    preset_id = uuid.uuid4()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="agent_preset",
        entity_id=preset_id,
        agent_preset_id=preset_id,
        agent_preset_version_id=uuid.uuid4(),
        agents_binding={"enabled": True, "subagents": []},
    )
    validate_mock = AsyncMock()
    agents_binding_mock = AsyncMock(return_value=None)
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(agent_preset_version_id=None),
    )

    validate_mock.assert_not_awaited()
    assert updated.agent_preset_id == preset_id
    assert updated.agent_preset_version_id is None
    assert updated.agents_binding is None
    agents_binding_mock.assert_awaited_once_with(None)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_update_session_ignores_mismatched_preset_id_for_preset_sessions() -> (
    None
):
    service, session, role = _build_service()
    preset_id = uuid.uuid4()
    mismatched_preset_id = uuid.uuid4()
    new_version_id = uuid.uuid4()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="agent_preset",
        entity_id=preset_id,
        agent_preset_id=preset_id,
        agent_preset_version_id=uuid.uuid4(),
    )
    validate_mock = AsyncMock(return_value=new_version_id)
    agents_binding_mock = AsyncMock(return_value={"enabled": True, "subagents": []})
    service._validate_preset_version_for_assignment = validate_mock
    service._resolve_agents_binding_for_preset_version_id = agents_binding_mock

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(
            agent_preset_id=mismatched_preset_id,
            agent_preset_version_id=new_version_id,
        ),
    )

    validate_mock.assert_awaited_once_with(
        entity_type=AgentSessionEntity.AGENT_PRESET,
        entity_id=preset_id,
        agent_preset_id=mismatched_preset_id,
        agent_preset_version_id=new_version_id,
    )
    assert updated.agent_preset_id == preset_id
    assert updated.agent_preset_version_id == new_version_id
    assert updated.agents_binding == {"enabled": True, "subagents": []}
    agents_binding_mock.assert_awaited_once_with(new_version_id)
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_update_session_skips_entity_type_parsing_for_unrelated_updates() -> None:
    service, session, role = _build_service()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="not-a-real-entity",
        entity_id=uuid.uuid4(),
    )

    updated = await service.update_session(
        agent_session,
        params=AgentSessionUpdate(title="Renamed chat"),
    )

    assert updated.title == "Renamed chat"
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(agent_session)


@pytest.mark.anyio
async def test_update_session_rejects_preset_updates_for_invalid_entity_type() -> None:
    service, session, role = _build_service()
    agent_session = AgentSession(
        workspace_id=role.workspace_id,
        title="Chat",
        created_by=uuid.uuid4(),
        entity_type="not-a-real-entity",
        entity_id=uuid.uuid4(),
    )

    with pytest.raises(
        TracecatValidationError,
        match="Cannot update preset assignment for a session with an invalid entity type",
    ):
        await service.update_session(
            agent_session,
            params=AgentSessionUpdate(agent_preset_version_id=uuid.uuid4()),
        )

    session.commit.assert_not_awaited()
    session.refresh.assert_not_awaited()
