"""
Shared fixtures for all manager tests.

This conftest.py makes fixtures available to all test files in the tests/managers/ directory.
"""

import os
import time
import uuid
from typing import Tuple

import pytest
from anthropic.types.beta import BetaMessage
from anthropic.types.beta.messages import BetaMessageBatchIndividualResponse, BetaMessageBatchSucceededResult
from sqlalchemy import text

from letta.config import LettaConfig
from letta.functions.functions import derive_openai_json_schema, parse_source_code
from letta.functions.mcp_client.types import MCPTool
from letta.helpers import ToolRulesSolver
from letta.orm import Base
from letta.schemas.agent import CreateAgent
from letta.schemas.block import Block as PydanticBlock, CreateBlock
from letta.schemas.embedding_config import EmbeddingConfig
from letta.schemas.enums import JobStatus, MessageRole, RunStatus
from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate, SandboxEnvironmentVariableUpdate
from letta.schemas.file import FileMetadata as PydanticFileMetadata
from letta.schemas.job import BatchJob, Job as PydanticJob
from letta.schemas.letta_message_content import TextContent
from letta.schemas.llm_batch_job import AgentStepState
from letta.schemas.llm_config import LLMConfig
from letta.schemas.message import Message as PydanticMessage, MessageCreate
from letta.schemas.organization import Organization
from letta.schemas.passage import Passage as PydanticPassage
from letta.schemas.run import Run as PydanticRun
from letta.schemas.sandbox_config import E2BSandboxConfig, SandboxConfigCreate
from letta.schemas.source import Source as PydanticSource
from letta.schemas.tool import Tool as PydanticTool, ToolCreate
from letta.schemas.tool_rule import InitToolRule
from letta.schemas.user import User as PydanticUser
from letta.server.db import db_registry
from letta.server.server import SyncServer
from letta.services.block_manager import BlockManager

# Constants
DEFAULT_EMBEDDING_CONFIG = EmbeddingConfig.default_config(provider="openai")
CREATE_DELAY_SQLITE = 1
USING_SQLITE = not bool(os.getenv("LETTA_PG_URI"))


# ======================================================================================================================
# Database and Server Fixtures
# ======================================================================================================================


@pytest.fixture
async def async_session():
    """Provide an async database session."""
    async with db_registry.async_session() as session:
        yield session


@pytest.fixture(autouse=True)
async def _clear_tables(async_session):
    """Clear all tables before each test (except block_history)."""
    # Temporarily disable foreign key constraints for SQLite only
    engine_name = async_session.bind.dialect.name
    if engine_name == "sqlite":
        await async_session.execute(text("PRAGMA foreign_keys = OFF"))

    for table in reversed(Base.metadata.sorted_tables):  # Reverse to avoid FK issues
        # If this is the block_history table, skip it
        if table.name == "block_history":
            continue
        await async_session.execute(table.delete())  # Truncate table
    await async_session.commit()

    # Re-enable foreign key constraints for SQLite only
    if engine_name == "sqlite":
        await async_session.execute(text("PRAGMA foreign_keys = ON"))


@pytest.fixture(scope="module")
def server():
    """Create a server instance for the test module."""
    config = LettaConfig.load()
    config.save()
    server = SyncServer(init_with_default_org_and_user=False)
    return server


# ======================================================================================================================
# Organization and User Fixtures
# ======================================================================================================================


@pytest.fixture
async def default_organization(server: SyncServer):
    """Create and return the default organization."""
    org = await server.organization_manager.create_default_organization_async()
    yield org


@pytest.fixture
async def other_organization(server: SyncServer):
    """Create and return another organization."""
    org = await server.organization_manager.create_organization_async(pydantic_org=Organization(name="letta"))
    yield org


@pytest.fixture
async def default_user(server: SyncServer, default_organization):
    """Create and return the default user within the default organization."""
    user = await server.user_manager.create_default_actor_async(org_id=default_organization.id)
    yield user


@pytest.fixture
async def other_user(server: SyncServer, default_organization):
    """Create and return another user within the default organization."""
    user = await server.user_manager.create_actor_async(PydanticUser(name="other", organization_id=default_organization.id))
    yield user


@pytest.fixture
async def other_user_different_org(server: SyncServer, other_organization):
    """Create and return a user in a different organization."""
    user = await server.user_manager.create_actor_async(PydanticUser(name="other", organization_id=other_organization.id))
    yield user


# ======================================================================================================================
# Source and File Fixtures
# ======================================================================================================================


@pytest.fixture
async def default_source(server: SyncServer, default_user):
    """Create and return the default source."""
    source_pydantic = PydanticSource(
        name="Test Source",
        description="This is a test source.",
        metadata={"type": "test"},
        embedding_config=DEFAULT_EMBEDDING_CONFIG,
    )
    source = await server.source_manager.create_source(source=source_pydantic, actor=default_user)
    yield source


@pytest.fixture
async def other_source(server: SyncServer, default_user):
    """Create and return another source."""
    source_pydantic = PydanticSource(
        name="Another Test Source",
        description="This is yet another test source.",
        metadata={"type": "another_test"},
        embedding_config=DEFAULT_EMBEDDING_CONFIG,
    )
    source = await server.source_manager.create_source(source=source_pydantic, actor=default_user)
    yield source


@pytest.fixture
async def default_file(server: SyncServer, default_source, default_user, default_organization):
    """Create and return the default file."""
    file = await server.file_manager.create_file(
        PydanticFileMetadata(file_name="test_file", organization_id=default_organization.id, source_id=default_source.id),
        actor=default_user,
    )
    yield file


@pytest.fixture
async def another_file(server: SyncServer, default_source, default_user, default_organization):
    """Create and return another file."""
    pf = PydanticFileMetadata(
        file_name="another_file",
        organization_id=default_organization.id,
        source_id=default_source.id,
    )
    file = await server.file_manager.create_file(pf, actor=default_user)
    yield file


# ======================================================================================================================
# Tool Fixtures
# ======================================================================================================================


@pytest.fixture
async def print_tool(server: SyncServer, default_user, default_organization):
    """Create and return a print tool."""

    def print_tool(message: str):
        """
        Args:
            message (str): The message to print.

        Returns:
            str: The message that was printed.
        """
        print(message)
        return message

    # Set up tool details
    source_code = parse_source_code(print_tool)
    source_type = "python"
    description = "test_description"
    tags = ["test"]
    metadata = {"a": "b"}

    tool = PydanticTool(description=description, tags=tags, source_code=source_code, source_type=source_type, metadata_=metadata)
    derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name)

    derived_name = derived_json_schema["name"]
    tool.json_schema = derived_json_schema
    tool.name = derived_name

    tool = await server.tool_manager.create_or_update_tool_async(tool, actor=default_user)
    yield tool


@pytest.fixture
async def bash_tool(server: SyncServer, default_user, default_organization):
    """Create and return a bash tool with requires_approval."""

    def bash_tool(operation: str):
        """
        Args:
            operation (str): The bash operation to execute.

        Returns:
            str: The result of the executed operation.
        """
        print("scary bash operation")
        return "success"

    # Set up tool details
    source_code = parse_source_code(bash_tool)
    source_type = "python"
    description = "test_description"
    tags = ["test"]
    metadata = {"a": "b"}

    tool = PydanticTool(description=description, tags=tags, source_code=source_code, source_type=source_type, metadata_=metadata)
    derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name)

    derived_name = derived_json_schema["name"]
    tool.json_schema = derived_json_schema
    tool.name = derived_name
    tool.default_requires_approval = True

    tool = await server.tool_manager.create_or_update_tool_async(tool, actor=default_user)
    yield tool


@pytest.fixture
async def other_tool(server: SyncServer, default_user, default_organization):
    """Create and return another tool."""

    def print_other_tool(message: str):
        """
        Args:
            message (str): The message to print.

        Returns:
            str: The message that was printed.
        """
        print(message)
        return message

    # Set up tool details
    source_code = parse_source_code(print_other_tool)
    source_type = "python"
    description = "other_tool_description"
    tags = ["test"]

    tool = PydanticTool(description=description, tags=tags, source_code=source_code, source_type=source_type)
    derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name)

    derived_name = derived_json_schema["name"]
    tool.json_schema = derived_json_schema
    tool.name = derived_name

    tool = await server.tool_manager.create_or_update_tool_async(tool, actor=default_user)
    yield tool


# ======================================================================================================================
# Block Fixtures
# ======================================================================================================================


@pytest.fixture
async def default_block(server: SyncServer, default_user):
    """Create and return a default block."""
    block_manager = BlockManager()
    block_data = PydanticBlock(
        label="default_label",
        value="Default Block Content",
        description="A default test block",
        limit=1000,
        metadata={"type": "test"},
    )
    block = await block_manager.create_or_update_block_async(block_data, actor=default_user)
    yield block


@pytest.fixture
async def other_block(server: SyncServer, default_user):
    """Create and return another block."""
    block_manager = BlockManager()
    block_data = PydanticBlock(
        label="other_label",
        value="Other Block Content",
        description="Another test block",
        limit=500,
        metadata={"type": "test"},
    )
    block = await block_manager.create_or_update_block_async(block_data, actor=default_user)
    yield block


# ======================================================================================================================
# Agent Fixtures
# ======================================================================================================================


@pytest.fixture
async def sarah_agent(server: SyncServer, default_user, default_organization):
    """Create and return a sample agent named 'sarah_agent'."""
    agent_state = await server.agent_manager.create_agent_async(
        agent_create=CreateAgent(
            name="sarah_agent",
            agent_type="memgpt_v2_agent",
            memory_blocks=[],
            llm_config=LLMConfig.default_config("gpt-4o-mini"),
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            include_base_tools=False,
        ),
        actor=default_user,
    )
    yield agent_state


@pytest.fixture
async def charles_agent(server: SyncServer, default_user, default_organization):
    """Create and return a sample agent named 'charles_agent'."""
    agent_state = await server.agent_manager.create_agent_async(
        agent_create=CreateAgent(
            name="charles_agent",
            agent_type="memgpt_v2_agent",
            memory_blocks=[CreateBlock(label="human", value="Charles"), CreateBlock(label="persona", value="I am a helpful assistant")],
            llm_config=LLMConfig.default_config("gpt-4o-mini"),
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            include_base_tools=False,
        ),
        actor=default_user,
    )
    yield agent_state


@pytest.fixture
async def comprehensive_test_agent_fixture(server: SyncServer, default_user, print_tool, default_source, default_block) -> Tuple:
    """Create a comprehensive test agent with all features."""
    memory_blocks = [CreateBlock(label="human", value="BananaBoy"), CreateBlock(label="persona", value="I am a helpful assistant")]
    create_agent_request = CreateAgent(
        agent_type="memgpt_v2_agent",
        system="test system",
        memory_blocks=memory_blocks,
        llm_config=LLMConfig.default_config("gpt-4o-mini"),
        embedding_config=EmbeddingConfig.default_config(provider="openai"),
        block_ids=[default_block.id],
        tool_ids=[print_tool.id],
        source_ids=[default_source.id],
        tags=["a", "b"],
        description="test_description",
        metadata={"test_key": "test_value"},
        tool_rules=[InitToolRule(tool_name=print_tool.name)],
        initial_message_sequence=[MessageCreate(role=MessageRole.user, content="hello world")],
        tool_exec_environment_variables={"test_env_var_key_a": "test_env_var_value_a", "test_env_var_key_b": "test_env_var_value_b"},
        message_buffer_autoclear=True,
        include_base_tools=False,
    )
    created_agent = await server.agent_manager.create_agent_async(
        create_agent_request,
        actor=default_user,
    )

    yield created_agent, create_agent_request


# ======================================================================================================================
# Archive and Passage Fixtures
# ======================================================================================================================


@pytest.fixture
async def default_archive(server: SyncServer, default_user):
    """Create and return a default archive."""
    archive = await server.archive_manager.create_archive_async(
        "test", embedding_config=EmbeddingConfig.default_config(provider="openai"), actor=default_user
    )
    yield archive


@pytest.fixture
async def agent_passage_fixture(server: SyncServer, default_user, sarah_agent):
    """Create an agent passage."""
    # Get or create default archive for the agent
    archive = await server.archive_manager.get_or_create_default_archive_for_agent_async(agent_state=sarah_agent, actor=default_user)

    passage = await server.passage_manager.create_agent_passage_async(
        PydanticPassage(
            text="Hello, I am an agent passage",
            archive_id=archive.id,
            organization_id=default_user.organization_id,
            embedding=[0.1],
            embedding_config=DEFAULT_EMBEDDING_CONFIG,
            metadata={"type": "test"},
        ),
        actor=default_user,
    )
    yield passage


@pytest.fixture
async def source_passage_fixture(server: SyncServer, default_user, default_file, default_source):
    """Create a source passage."""
    passage = await server.passage_manager.create_source_passage_async(
        PydanticPassage(
            text="Hello, I am a source passage",
            source_id=default_source.id,
            file_id=default_file.id,
            organization_id=default_user.organization_id,
            embedding=[0.1],
            embedding_config=DEFAULT_EMBEDDING_CONFIG,
            metadata={"type": "test"},
        ),
        file_metadata=default_file,
        actor=default_user,
    )
    yield passage


# ======================================================================================================================
# Message Fixtures
# ======================================================================================================================


@pytest.fixture
async def hello_world_message_fixture(server: SyncServer, default_user, sarah_agent):
    """Create a hello world message."""
    message = PydanticMessage(
        agent_id=sarah_agent.id,
        role="user",
        content=[TextContent(text="Hello, world!")],
    )

    msg = await server.message_manager.create_many_messages_async([message], actor=default_user)
    yield msg[0]


# ======================================================================================================================
# Sandbox Fixtures
# ======================================================================================================================


@pytest.fixture
async def sandbox_config_fixture(server: SyncServer, default_user):
    """Create a sandbox configuration."""
    sandbox_config_create = SandboxConfigCreate(
        config=E2BSandboxConfig(),
    )
    created_config = await server.sandbox_config_manager.create_or_update_sandbox_config_async(sandbox_config_create, actor=default_user)
    yield created_config


@pytest.fixture
async def sandbox_env_var_fixture(server: SyncServer, sandbox_config_fixture, default_user):
    """Create a sandbox environment variable."""
    env_var_create = SandboxEnvironmentVariableCreate(
        key="SAMPLE_VAR",
        value="sample_value",
        description="A sample environment variable for testing.",
    )
    created_env_var = await server.sandbox_config_manager.create_sandbox_env_var_async(
        env_var_create, sandbox_config_id=sandbox_config_fixture.id, actor=default_user
    )
    yield created_env_var


# ======================================================================================================================
# File Attachment Fixtures
# ======================================================================================================================


@pytest.fixture
async def file_attachment(server: SyncServer, default_user, sarah_agent, default_file):
    """Create a file attachment to an agent."""
    assoc, closed_files = await server.file_agent_manager.attach_file(
        agent_id=sarah_agent.id,
        file_id=default_file.id,
        file_name=default_file.file_name,
        source_id=default_file.source_id,
        actor=default_user,
        visible_content="initial",
        max_files_open=sarah_agent.max_files_open,
    )
    yield assoc


# ======================================================================================================================
# Job Fixtures
# ======================================================================================================================


@pytest.fixture
async def default_job(server: SyncServer, default_user):
    """Create and return a default job."""
    job_pydantic = PydanticJob(
        user_id=default_user.id,
        status=JobStatus.pending,
    )
    job = await server.job_manager.create_job_async(pydantic_job=job_pydantic, actor=default_user)
    yield job


@pytest.fixture
async def default_run(server: SyncServer, default_user, sarah_agent):
    """Create and return a default run."""
    run_pydantic = PydanticRun(
        agent_id=sarah_agent.id,
        status=RunStatus.created,
    )
    run = await server.run_manager.create_run(pydantic_run=run_pydantic, actor=default_user)
    yield run


@pytest.fixture
async def letta_batch_job(server: SyncServer, default_user):
    """Create a batch job."""
    return await server.job_manager.create_job_async(BatchJob(user_id=default_user.id), actor=default_user)


# ======================================================================================================================
# MCP Tool Fixtures
# ======================================================================================================================


@pytest.fixture
async def mcp_tool(server: SyncServer, default_user):
    """Create an MCP tool."""
    mcp_tool = MCPTool(
        name="weather_lookup",
        description="Fetches the current weather for a given location.",
        inputSchema={
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "The name of the city or location."},
                "units": {
                    "type": "string",
                    "enum": ["metric", "imperial"],
                    "description": "The unit system for temperature (metric or imperial).",
                },
            },
            "required": ["location"],
        },
    )
    mcp_server_name = "test"
    mcp_server_id = "test-server-id"  # Mock server ID for testing
    tool_create = ToolCreate.from_mcp(mcp_server_name=mcp_server_name, mcp_tool=mcp_tool)
    tool = await server.tool_manager.create_or_update_mcp_tool_async(
        tool_create=tool_create, mcp_server_name=mcp_server_name, mcp_server_id=mcp_server_id, actor=default_user
    )
    yield tool


# ======================================================================================================================
# Test Data Creation Fixtures
# ======================================================================================================================


@pytest.fixture
async def create_test_passages(server: SyncServer, default_file, default_user, sarah_agent, default_source):
    """Helper function to create test passages for all tests."""
    # Get or create default archive for the agent
    archive = await server.archive_manager.get_or_create_default_archive_for_agent(
        agent_id=sarah_agent.id, agent_name=sarah_agent.name, actor=default_user
    )

    # Create agent passages
    passages = []
    for i in range(5):
        passage = await server.passage_manager.create_agent_passage(
            PydanticPassage(
                text=f"Agent passage {i}",
                archive_id=archive.id,
                organization_id=default_user.organization_id,
                embedding=[0.1],
                embedding_config=DEFAULT_EMBEDDING_CONFIG,
                metadata={"type": "test"},
            ),
            actor=default_user,
        )
        passages.append(passage)
        if USING_SQLITE:
            time.sleep(CREATE_DELAY_SQLITE)

    # Create source passages
    for i in range(5):
        passage = await server.passage_manager.create_source_passage(
            PydanticPassage(
                text=f"Source passage {i}",
                source_id=default_source.id,
                file_id=default_file.id,
                organization_id=default_user.organization_id,
                embedding=[0.1],
                embedding_config=DEFAULT_EMBEDDING_CONFIG,
                metadata={"type": "test"},
            ),
            file_metadata=default_file,
            actor=default_user,
        )
        passages.append(passage)
        if USING_SQLITE:
            time.sleep(CREATE_DELAY_SQLITE)

    return passages


@pytest.fixture
async def agent_passages_setup(server: SyncServer, default_archive, default_source, default_file, default_user, sarah_agent):
    """Setup fixture for agent passages tests."""
    agent_id = sarah_agent.id
    actor = default_user

    await server.agent_manager.attach_source_async(agent_id=agent_id, source_id=default_source.id, actor=actor)

    # Create some source passages
    source_passages = []
    for i in range(3):
        passage = await server.passage_manager.create_source_passage_async(
            PydanticPassage(
                organization_id=actor.organization_id,
                source_id=default_source.id,
                file_id=default_file.id,
                text=f"Source passage {i}",
                embedding=[0.1],  # Default OpenAI embedding size
                embedding_config=DEFAULT_EMBEDDING_CONFIG,
            ),
            file_metadata=default_file,
            actor=actor,
        )
        source_passages.append(passage)

    # attach archive
    await server.archive_manager.attach_agent_to_archive_async(
        agent_id=agent_id, archive_id=default_archive.id, is_owner=True, actor=default_user
    )

    # Create some agent passages
    agent_passages = []
    for i in range(2):
        passage = await server.passage_manager.create_agent_passage_async(
            PydanticPassage(
                organization_id=actor.organization_id,
                archive_id=default_archive.id,
                text=f"Agent passage {i}",
                embedding=[0.1],  # Default OpenAI embedding size
                embedding_config=DEFAULT_EMBEDDING_CONFIG,
            ),
            actor=actor,
        )
        agent_passages.append(passage)

    yield agent_passages, source_passages

    # Cleanup
    await server.source_manager.delete_source(default_source.id, actor=actor)


@pytest.fixture
async def agent_with_tags(server: SyncServer, default_user):
    """Create agents with specific tags."""
    agent1 = await server.agent_manager.create_agent_async(
        agent_create=CreateAgent(
            name="agent1",
            agent_type="memgpt_v2_agent",
            tags=["primary_agent", "benefit_1"],
            llm_config=LLMConfig.default_config("gpt-4o-mini"),
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            memory_blocks=[],
            include_base_tools=False,
        ),
        actor=default_user,
    )

    agent2 = await server.agent_manager.create_agent_async(
        agent_create=CreateAgent(
            name="agent2",
            agent_type="memgpt_v2_agent",
            tags=["primary_agent", "benefit_2"],
            llm_config=LLMConfig.default_config("gpt-4o-mini"),
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            memory_blocks=[],
            include_base_tools=False,
        ),
        actor=default_user,
    )

    agent3 = await server.agent_manager.create_agent_async(
        agent_create=CreateAgent(
            name="agent3",
            agent_type="memgpt_v2_agent",
            tags=["primary_agent", "benefit_1", "benefit_2"],
            llm_config=LLMConfig.default_config("gpt-4o-mini"),
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            memory_blocks=[],
            include_base_tools=False,
        ),
        actor=default_user,
    )

    return [agent1, agent2, agent3]


# ======================================================================================================================
# LLM and Step State Fixtures
# ======================================================================================================================


@pytest.fixture
def dummy_llm_config() -> LLMConfig:
    """Create a dummy LLM config."""
    return LLMConfig.default_config("gpt-4o-mini")


@pytest.fixture
def dummy_tool_rules_solver() -> ToolRulesSolver:
    """Create a dummy tool rules solver."""
    return ToolRulesSolver(tool_rules=[InitToolRule(tool_name="send_message")])


@pytest.fixture
def dummy_step_state(dummy_tool_rules_solver: ToolRulesSolver) -> AgentStepState:
    """Create a dummy step state."""
    return AgentStepState(step_number=1, tool_rules_solver=dummy_tool_rules_solver)


@pytest.fixture
def dummy_successful_response() -> BetaMessageBatchIndividualResponse:
    """Create a dummy successful Anthropic message batch response."""
    return BetaMessageBatchIndividualResponse(
        custom_id="my-second-request",
        result=BetaMessageBatchSucceededResult(
            type="succeeded",
            message=BetaMessage(
                id="msg_abc123",
                role="assistant",
                type="message",
                model="claude-3-5-sonnet-20240620",
                content=[{"type": "text", "text": "hi!"}],
                usage={"input_tokens": 5, "output_tokens": 7},
                stop_reason="end_turn",
            ),
        ),
    )


# ======================================================================================================================
# Environment Setup Fixtures
# ======================================================================================================================


@pytest.fixture(params=[None, "prod"])
def set_letta_environment(request, monkeypatch):
    """Parametrized fixture to test with different environment settings."""
    from letta.settings import settings

    # Patch the settings.environment attribute
    original = settings.environment
    monkeypatch.setattr(settings, "environment", request.param)
    yield request.param
    # Restore original environment
    monkeypatch.setattr(settings, "environment", original)