"""Tests for DBOS workflow functionality."""

import pytest
import pytest_asyncio
import os
import asyncio
import uuid
from timestep import (
    configure_dbos,
    ensure_dbos_launched,
    cleanup_dbos,
    run_agent_workflow,
    queue_agent_workflow,
    create_scheduled_agent_workflow,
    register_generic_workflows,
)
from timestep._vendored_imports import (
    Agent,
    OpenAIConversationsSession,
    ModelSettings,
    TResponseInputItem,
)
from timestep.stores.agent_store.store import save_agent
from timestep.stores.session_store.store import save_session


@pytest_asyncio.fixture(scope="function")
async def setup_dbos():
    """Set up DBOS for testing - follows DBOS recommended pattern."""
    from dbos import DBOS, DBOSConfig
    from timestep.config.dbos_config import _dbos_context
    import uuid
    
    # Step 1: Destroy DBOS (per DBOS testing best practices)
    # Note: Don't destroy registry - workflows are registered via decorators at import time
    try:
        DBOS.destroy(destroy_registry=False)
    except Exception:
        pass  # Ignore if DBOS doesn't exist yet
    
    # Step 2: Configure DBOS directly (matching DBOS recommended pattern)
    # Get connection string from environment
    db_url = os.environ.get("PG_CONNECTION_URI")
    if not db_url:
        raise ValueError("PG_CONNECTION_URI not set. Run 'make test-setup' to start the test database.")
    
    config: DBOSConfig = {
        "name": f"timestep-test-{uuid.uuid4().hex[:8]}",
        "system_database_url": db_url,
    }
    DBOS(config=config)
    
    # Update our context so get_dbos_connection_string() works
    _dbos_context.set_config(config)
    _dbos_context.set_configured(True)
    
    # Step 3: Reset system database (per DBOS testing best practices)
    try:
        DBOS.reset_system_database()
    except Exception:
        pass  # Ignore if reset fails (e.g., database in use)
    
    # Step 4: Launch DBOS (workflows are already registered via decorators)
    DBOS.launch()
    _dbos_context.set_launched(True)
    
    yield
    
    # Cleanup: shutdown and destroy
    try:
        if hasattr(DBOS, 'shutdown'):
            DBOS.shutdown()
        else:
            DBOS.destroy(destroy_registry=False)
    except Exception:
        pass
    finally:
        _dbos_context.set_configured(False)
        _dbos_context.set_launched(False)
        _dbos_context.set_config(None)


@pytest.mark.asyncio
async def test_configure_dbos(setup_dbos):
    """Test that DBOS can be configured."""
    # Configuration is done in fixture
    assert True  # If we get here, configuration worked


@pytest.mark.asyncio
@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set"
)
@pytest.mark.parametrize("model", ["gpt-4.1", "ollama/gpt-oss:20b-cloud", "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M"])
async def test_run_agent_workflow_basic(setup_dbos, model):
    """Test basic durable workflow execution."""
    if model == "ollama/gpt-oss:20b-cloud" or model == "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M":
        # Expected failure: Ollama cloud model has known compatibility issues
        with pytest.raises(Exception):
            # Create agent and session
            import uuid
            agent = Agent(
                instructions="You are a helpful assistant. Answer concisely.",
                model=model,
                model_settings=ModelSettings(temperature=0),
                name=f"Test Assistant Basic {uuid.uuid4().hex[:8]}",
            )
            session = OpenAIConversationsSession()
            
            # Save agent and session to database (stores manage connections internally)
            agent_id = await save_agent(agent)
            session_id = await save_session(session)
            
            input_items: list[TResponseInputItem] = [
                {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say 'hello' and nothing else."}]}
            ]
            
            await run_agent_workflow(
                agent_id=agent_id,
                input_items=input_items,
                session_id=session_id,
                stream=False,
                workflow_id="test-workflow-1"
            )
        return
    
    # Create agent and session
    import uuid
    agent = Agent(
        instructions="You are a helpful assistant. Answer concisely.",
        model=model,
        model_settings=ModelSettings(temperature=0),
        name=f"Test Assistant Basic {uuid.uuid4().hex[:8]}",
    )
    session = OpenAIConversationsSession()
    
    # Save agent and session to database (stores manage connections internally)
    agent_id = await save_agent(agent)
    session_id = await save_session(session)
    
    input_items: list[TResponseInputItem] = [
        {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say 'hello' and nothing else."}]}
    ]
    
    result = await run_agent_workflow(
        agent_id=agent_id,
        input_items=input_items,
        session_id=session_id,
        stream=False,
        workflow_id="test-workflow-1"
    )
    
    assert result is not None
    # Result is now a dict, not an object
    assert 'output' in result
    assert result['output'] is not None, f"Result output is None. Full result: {result}"
    
    # Extract text from output array if needed
    output_text = result['output']
    if isinstance(output_text, list):
        # Extract text from message items
        text_parts = []
        for item in output_text:
            if item.get('type') == 'message' and item.get('role') == 'assistant' and item.get('content'):
                for block in item['content']:
                    if block.get('type') == 'output_text' and block.get('text'):
                        text_parts.append(block['text'])
        output_text = ' '.join(text_parts)
    assert "hello" in str(output_text).lower(), f"Output text '{output_text}' does not contain 'hello'"


@pytest.mark.asyncio
@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set"
)
@pytest.mark.parametrize("model", ["gpt-4.1", "ollama/gpt-oss:20b-cloud", "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M"])
async def test_queue_agent_workflow(setup_dbos, model):
    """Test queued workflow execution."""
    if model == "ollama/gpt-oss:20b-cloud" or model == "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M":
        # Expected failure: Ollama cloud model has known compatibility issues (may work intermittently)
        try:
            # Create agent and session
            import uuid
            agent = Agent(
                instructions="You are a helpful assistant. Answer concisely.",
                model=model,
                model_settings=ModelSettings(temperature=0),
                name=f"Test Assistant Queue {uuid.uuid4().hex[:8]}",
            )
            session = OpenAIConversationsSession()
            
            # Save agent and session to database (stores manage connections internally)
            agent_id = await save_agent(agent)
            session_id = await save_session(session)
            
            input_items: list[TResponseInputItem] = [
                {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say 'queued' and nothing else."}]}
            ]
            
            dedup_id = f"test-queue-{uuid.uuid4().hex[:8]}"
            handle = await queue_agent_workflow(
                agent_id=agent_id,
                input_items=input_items,
                session_id=session_id,
                stream=False,
                deduplication_id=dedup_id
            )
            
            # Poll for workflow completion before getting result
            import time
            max_wait = 90
            start_time = time.time()
            status = None
            last_status = None
            
            while time.time() - start_time < max_wait:
                status_obj = handle.get_status()
                status = status_obj.status if hasattr(status_obj, 'status') else str(status_obj)
                
                if status != last_status:
                    elapsed = time.time() - start_time
                    print(f"[{elapsed:.1f}s] Status: {last_status} -> {status}")
                    last_status = status
                
                if str(status) in ['SUCCESS', 'FAILED', 'ERROR']:
                    break
                await asyncio.sleep(1)
            
            if status is None or str(status) not in ['SUCCESS', 'FAILED', 'ERROR']:
                pytest.fail(f"Workflow did not complete after {max_wait} seconds. Status: {status}")
            
            # Try to get result - if it throws, that's expected for these models
            try:
                result = handle.get_result()
                assert result is not None
                assert 'output' in result
                # If it works, that's acceptable (test may work intermittently)
            except Exception:
                # Expected to fail when getting result - this is the known failure case
                pass
        except Exception:
            # Expected to fail - this is the known failure case
            pass
        return
    
    # Create agent and session
    import uuid
    agent = Agent(
        instructions="You are a helpful assistant. Answer concisely.",
        model=model,
        model_settings=ModelSettings(temperature=0),
        name=f"Test Assistant Queue {uuid.uuid4().hex[:8]}",
    )
    session = OpenAIConversationsSession()
    
    # Save agent and session to database (stores manage connections internally)
    agent_id = await save_agent(agent)
    session_id = await save_session(session)
    
    input_items: list[TResponseInputItem] = [
        {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say 'queued' and nothing else."}]}
    ]
    
    dedup_id = f"test-queue-{uuid.uuid4().hex[:8]}"
    handle = await queue_agent_workflow(
        agent_id=agent_id,
        input_items=input_items,
        session_id=session_id,
        stream=False,
        deduplication_id=dedup_id
    )
    
    # Poll for workflow completion before getting result
    # Use same pattern as the working simplified test
    import time
    max_wait = 90  # Increase timeout
    start_time = time.time()
    status = None
    last_status = None
    
    while time.time() - start_time < max_wait:
        status_obj = handle.get_status()
        status = status_obj.status if hasattr(status_obj, 'status') else str(status_obj)
        
        # Log status changes for debugging
        if status != last_status:
            elapsed = time.time() - start_time
            print(f"[{elapsed:.1f}s] Status: {last_status} -> {status}")
            last_status = status
        
        if str(status) in ['SUCCESS', 'FAILED', 'ERROR']:
            break
        await asyncio.sleep(1)  # Check every second instead of 0.5s
    
    if status is None or str(status) not in ['SUCCESS', 'FAILED', 'ERROR']:
        pytest.fail(f"Workflow did not complete after {max_wait} seconds. Status: {status}")
    
    # get_result() is synchronous and blocks until workflow completes
    result = handle.get_result()
    assert result is not None
    # Result is now a dict, not an object
    assert 'output' in result
    assert result['output'] is not None, f"Result output is None. Full result: {result}"
    
    # Extract text from output array if needed
    output_text = result['output']
    if isinstance(output_text, list):
        # Extract text from message items
        text_parts = []
        for item in output_text:
            if item.get('type') == 'message' and item.get('role') == 'assistant' and item.get('content'):
                for block in item['content']:
                    if block.get('type') == 'output_text' and block.get('text'):
                        text_parts.append(block['text'])
        output_text = ' '.join(text_parts)
    assert "queued" in str(output_text).lower(), f"Output text '{output_text}' does not contain 'queued'"


@pytest.mark.asyncio
@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set"
)
@pytest.mark.parametrize("model", ["gpt-4.1", "ollama/gpt-oss:20b-cloud", "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M"])
async def test_create_scheduled_workflow(setup_dbos, model):
    """Test that scheduled workflows must be created before DBOS launch."""
    # This test verifies that scheduled workflows must be created before DBOS launch.
    # Since DBOS is launched in the fixture, this test should fail with an appropriate error.
    import uuid
    agent = Agent(
        instructions="You are a helpful assistant.",
        model=model,
        model_settings=ModelSettings(temperature=0),
        name=f"Test Assistant Scheduled {uuid.uuid4().hex[:8]}",
    )
    session = OpenAIConversationsSession()
    
    # Save agent and session to database (stores manage connections internally)
    agent_id = await save_agent(agent)
    session_id = await save_session(session)
    
    input_items: list[TResponseInputItem] = [
        {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]}
    ]
    
    # This should raise an error because DBOS is already launched
    # For ollama model, we expect any exception (compatibility issues)
    if model == "ollama/gpt-oss:20b-cloud" or model == "ollama/hf.co/mjschock/SmolVLM2-500M-Video-Instruct-GGUF:Q4_K_M":
        with pytest.raises(Exception):
            await create_scheduled_agent_workflow(
                crontab="0 * * * *",  # Every hour
                agent_id=agent_id,
                input_items=input_items,
                session_id=session_id,
                stream=False
            )
    else:
        with pytest.raises(RuntimeError, match="Cannot create scheduled workflow after DBOS launch"):
            await create_scheduled_agent_workflow(
                crontab="0 * * * *",  # Every hour
                agent_id=agent_id,
                input_items=input_items,
                session_id=session_id,
                stream=False
            )

