"""Tests for AgentRunner run ledger pull API authorization."""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest
from sqlalchemy.ext.asyncio import create_async_engine

from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore
from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry
from langbot.pkg.entity.persistence import agent_run as agent_run_model
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
from langbot_plugin.api.entities.builtin.agent_runner.run_ledger import (
    AgentRun,
    AgentRunEvent,
    RunEventPage,
    RunPage,
)
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction

from .conftest import make_resources


class FakeConnection:
    pass


class FakeApplication:
    def __init__(self, db_engine):
        self.logger = MagicMock()
        self.persistence_mgr = MagicMock()
        self.persistence_mgr.get_db_engine = MagicMock(return_value=db_engine)


@pytest.fixture
def session_registry(monkeypatch):
    registry = AgentRunSessionRegistry()
    monkeypatch.setattr(
        'langbot.pkg.agent.runner.session_registry._global_registry',
        registry,
    )
    return registry


@pytest.fixture
async def db_engine():
    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
    assert agent_run_model.AgentRun.__tablename__ == 'agent_run'
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield engine
    await engine.dispose()


def _handler(db_engine):
    async def fake_disconnect():
        return True

    fake_app = FakeApplication(db_engine)
    return RuntimeConnectionHandler(FakeConnection(), fake_disconnect, fake_app)


async def _register_session(
    session_registry,
    *,
    run_id='run_1',
    conversation_id='conv_1',
    available_apis=None,
):
    await session_registry.register(
        run_id=run_id,
        runner_id='plugin:test/runner/default',
        query_id=None,
        plugin_identity='test/runner',
        resources=make_resources(),
        conversation_id=conversation_id,
        bot_id='bot_1',
        workspace_id='workspace_1',
        thread_id=None,
        available_apis=available_apis or {},
    )


async def _create_run(
    db_engine,
    *,
    run_id='run_1',
    conversation_id='conv_1',
    bot_id='bot_1',
    workspace_id='workspace_1',
    thread_id=None,
    plugin_identity='test/runner',
    available_apis=None,
):
    store = RunLedgerStore(db_engine)
    await store.create_run(
        run_id=run_id,
        event_id='evt_1',
        binding_id='binding_1',
        runner_id='plugin:test/runner/default',
        conversation_id=conversation_id,
        bot_id=bot_id,
        workspace_id=workspace_id,
        thread_id=thread_id,
        authorization={
            'runner_id': 'plugin:test/runner/default',
            'binding_id': 'binding_1',
            'plugin_identity': plugin_identity,
            'resources': make_resources(),
            'available_apis': available_apis or {},
            'conversation_id': conversation_id,
            'bot_id': bot_id,
            'workspace_id': workspace_id,
            'thread_id': thread_id,
            'state_policy': {'enable_state': True, 'state_scopes': ['conversation', 'actor']},
            'state_context': {},
        },
    )
    await store.append_event(
        run_id=run_id,
        sequence=1,
        event_type='message.completed',
        data={'message': {'role': 'assistant', 'content': 'ok'}},
    )


@pytest.mark.asyncio
async def test_run_get_returns_current_run(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_get': True})
    await _create_run(db_engine)
    handler = _handler(db_engine)
    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]

    result = await run_get(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code == 0
    run = AgentRun.model_validate(result.data)
    assert run.run_id == 'run_1'
    assert run.status == 'running'


@pytest.mark.asyncio
async def test_run_list_rejects_cross_conversation(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_list': True})
    handler = _handler(db_engine)
    run_list = handler.actions[PluginToRuntimeAction.RUN_LIST.value]

    result = await run_list(
        {
            'run_id': 'run_1',
            'conversation_id': 'conv_other',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code != 0
    assert 'not accessible' in result.message.lower()


@pytest.mark.asyncio
async def test_run_list_returns_scoped_runs(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_list': True})
    await _create_run(db_engine)
    await _create_run(db_engine, run_id='run_other', conversation_id='conv_other')
    handler = _handler(db_engine)
    run_list = handler.actions[PluginToRuntimeAction.RUN_LIST.value]

    result = await run_list(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code == 0
    page = RunPage.model_validate(result.data)
    assert [run.run_id for run in page.items] == ['run_1']


@pytest.mark.asyncio
async def test_run_events_page_returns_events(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_events_page': True})
    await _create_run(db_engine)
    handler = _handler(db_engine)
    run_events_page = handler.actions[PluginToRuntimeAction.RUN_EVENTS_PAGE.value]

    result = await run_events_page(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code == 0
    page = RunEventPage.model_validate(result.data)
    assert [item.sequence for item in page.items] == [1]
    assert page.items[0].type == 'message.completed'


@pytest.mark.asyncio
async def test_run_get_uses_persistent_authorization_after_session_expired(db_engine):
    await _create_run(db_engine, available_apis={'run_get': True})
    handler = _handler(db_engine)
    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]

    result = await run_get(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code == 0
    run = AgentRun.model_validate(result.data)
    assert run.run_id == 'run_1'


@pytest.mark.asyncio
async def test_persistent_run_get_rejects_cross_scope(db_engine):
    await _create_run(db_engine, available_apis={'run_get': True})
    await _create_run(
        db_engine,
        run_id='run_other',
        conversation_id='conv_other',
        available_apis={'run_get': True},
    )
    handler = _handler(db_engine)
    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]

    result = await run_get(
        {
            'run_id': 'run_1',
            'target_run_id': 'run_other',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code != 0
    assert 'not accessible' in result.message.lower()


@pytest.mark.asyncio
async def test_persistent_run_get_requires_capability(db_engine):
    await _create_run(db_engine, available_apis={'run_get': False})
    handler = _handler(db_engine)
    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]

    result = await run_get(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code != 0
    assert 'not authorized' in result.message.lower()


@pytest.mark.asyncio
async def test_persistent_authorization_does_not_reopen_artifact_api(db_engine):
    await _create_run(db_engine, available_apis={'artifact_metadata': True})
    handler = _handler(db_engine)
    artifact_metadata = handler.actions[PluginToRuntimeAction.ARTIFACT_METADATA.value]

    result = await artifact_metadata(
        {
            'run_id': 'run_1',
            'artifact_id': 'artifact_1',
            'caller_plugin_identity': 'test/runner',
        }
    )

    assert result.code != 0
    assert 'not found or expired' in result.message.lower()


@pytest.mark.asyncio
async def test_run_cancel_basic_path(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_cancel': True})
    await _create_run(db_engine)
    handler = _handler(db_engine)
    run_cancel = handler.actions[PluginToRuntimeAction.RUN_CANCEL.value]

    result = await run_cancel(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'reason': 'user requested',
        }
    )

    assert result.code == 0
    run = AgentRun.model_validate(result.data)
    assert run.run_id == 'run_1'
    assert run.cancel_requested_at is not None
    assert run.status_reason == 'user requested'


@pytest.mark.asyncio
async def test_run_append_result_basic_path(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_append_result': True})
    await _create_run(db_engine)
    handler = _handler(db_engine)
    run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value]

    result = await run_append_result(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'sequence': 2,
            'result': {
                'type': 'message.delta',
                'data': {'delta': 'hello'},
                'usage': {'output_tokens': 1},
            },
        }
    )

    assert result.code == 0
    event = AgentRunEvent.model_validate(result.data)
    assert event.run_id == 'run_1'
    assert event.sequence == 2
    assert event.type == 'message.delta'
    assert event.data == {'delta': 'hello'}
    assert event.usage == {'output_tokens': 1}


@pytest.mark.asyncio
async def test_run_finalize_basic_path(session_registry, db_engine):
    await _register_session(session_registry, available_apis={'run_finalize': True})
    await _create_run(db_engine)
    handler = _handler(db_engine)
    run_finalize = handler.actions[PluginToRuntimeAction.RUN_FINALIZE.value]

    result = await run_finalize(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'status': 'completed',
            'status_reason': 'done',
            'usage': {'total_tokens': 3},
        }
    )

    assert result.code == 0
    run = AgentRun.model_validate(result.data)
    assert run.status == 'completed'
    assert run.status_reason == 'done'
    assert run.finished_at is not None
    assert run.usage == {'total_tokens': 3}


@pytest.mark.asyncio
async def test_runtime_register_heartbeat_and_list_actions(session_registry, db_engine):
    await _register_session(
        session_registry,
        available_apis={
            'runtime_register': True,
            'runtime_heartbeat': True,
            'runtime_list': True,
        },
    )
    handler = _handler(db_engine)
    runtime_register = handler.actions[PluginToRuntimeAction.RUNTIME_REGISTER.value]
    runtime_heartbeat = handler.actions[PluginToRuntimeAction.RUNTIME_HEARTBEAT.value]
    runtime_list = handler.actions[PluginToRuntimeAction.RUNTIME_LIST.value]

    registered = await runtime_register(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'runtime_id': 'runtime_1',
            'display_name': 'Runtime 1',
            'capabilities': {'runner': True},
            'labels': {'region': 'test'},
            'metadata': {'slots': 2},
        }
    )

    assert registered.code == 0
    assert registered.data['runtime_id'] == 'runtime_1'
    assert registered.data['capabilities'] == {'runner': True}

    heartbeat = await runtime_heartbeat(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'runtime_id': 'runtime_1',
            'capabilities': {'runner': True, 'stream': True},
            'labels': {'region': 'test'},
            'metadata': {'active_runs': 1},
        }
    )

    assert heartbeat.code == 0
    assert heartbeat.data['capabilities'] == {'runner': True, 'stream': True}
    assert heartbeat.data['metadata'] == {'slots': 2, 'active_runs': 1}

    page = await runtime_list(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'statuses': ['online'],
            'labels': {'region': 'test'},
        }
    )

    assert page.code == 0
    assert [item['runtime_id'] for item in page.data['items']] == ['runtime_1']


@pytest.mark.asyncio
async def test_run_claim_renew_and_release_actions(session_registry, db_engine):
    await _register_session(
        session_registry,
        available_apis={
            'run_claim': True,
            'run_renew_claim': True,
            'run_release_claim': True,
        },
    )
    await RunLedgerStore(db_engine).create_run(
        run_id='queued_run',
        event_id='evt_queued',
        binding_id='binding_1',
        runner_id='plugin:test/runner/default',
        conversation_id='conv_1',
        bot_id='bot_1',
        workspace_id='workspace_1',
        status='queued',
        queue_name='default',
        priority=5,
    )
    handler = _handler(db_engine)
    run_claim = handler.actions[PluginToRuntimeAction.RUN_CLAIM.value]
    run_renew_claim = handler.actions[PluginToRuntimeAction.RUN_RENEW_CLAIM.value]
    run_release_claim = handler.actions[PluginToRuntimeAction.RUN_RELEASE_CLAIM.value]

    claimed = await run_claim(
        {
            'run_id': 'run_1',
            'caller_plugin_identity': 'test/runner',
            'runtime_id': 'runtime_1',
            'queue_name': 'default',
            'lease_seconds': 30,
        }
    )

    assert claimed.code == 0
    assert claimed.data['run_id'] == 'queued_run'
    assert claimed.data['status'] == 'claimed'
    assert claimed.data['claimed_by_runtime_id'] == 'runtime_1'
    claim_token = claimed.data['claim_token']

    renewed = await run_renew_claim(
        {
            'run_id': 'run_1',
            'target_run_id': 'queued_run',
            'caller_plugin_identity': 'test/runner',
            'runtime_id': 'runtime_1',
            'claim_token': claim_token,
            'lease_seconds': 60,
        }
    )

    assert renewed.code == 0
    assert renewed.data['claim_token'] == claim_token

    released = await run_release_claim(
        {
            'run_id': 'run_1',
            'target_run_id': 'queued_run',
            'caller_plugin_identity': 'test/runner',
            'runtime_id': 'runtime_1',
            'claim_token': claim_token,
            'reason': 'done with lease',
        }
    )

    assert released.code == 0
    assert released.data['status'] == 'queued'
    assert released.data['claimed_by_runtime_id'] is None
    assert released.data['claim_token'] is None