import random

from ismcore.model.processor_state import (
    State,
    StateConfig,
    StateConfigLM,
    StateDataKeyDefinition,
    StateDataColumnDefinition)


def create_mock_state_json():
    state = create_mock_state()
    state_json = state.model_dump()
    return state_json, state


def create_mock_state_with_no_data(state_id: str = None) -> State:
    state = State(
        id=state_id,
        config=StateConfig(
            name="Test Me",
            storage_class="database"
        )
    )
    return state


def create_mock_state(state_id: str = None) -> State:
    state = State(
        id=state_id,
        config=StateConfigLM(
            name="Test Me",
            storage_class="database",
            user_template_id="/tmp/tmp_user_template",
            system_template_id="/tmp/tmp_system_template"
        )
    )

    for i in range(10):
        state.process_and_add_columns({
            "State Key": "hello_world_state_key_a",
            "State Data": random.randbytes(64)
        })

    return state

def test_state_has_query_state():
    state = State(
        config=StateConfig(
            name="Test State Hash Query State Key",
            primary_key=[
                StateDataKeyDefinition(name="question"),
                StateDataKeyDefinition(name="animal")
            ],
            template_columns=[
                StateDataKeyDefinition(name="question")
            ]
        )
    )

    query_state = state.apply_query_state(query_state={
        "question": "what do you think about {animal}s?",
        "animal": "cat"
    })

    query_state = state.apply_query_state(query_state={
        "question": "what do you think about {animal}s?",
        "animal": "dog"
    })

    query_state_0 = state.build_query_state_from_row_data(0)
    query_state_1 = state.build_query_state_from_row_data(1)

    assert query_state_0["question"] == "what do you think about cats?"
    assert query_state_1["question"] == "what do you think about dogs?"

    query_state = state.apply_query_state(query_state={
        "question": "what do you think about {animal}s?",
        "animal": "dog"
    })

    assert state.count == 2

    query_state = state.apply_query_state(query_state={
        "question": "what do you think about {animal}s?",
        "animal": "pig"
    })

    assert state.count == 3


def test_state_json_model():
    state_json, state = create_mock_state_json()
    deserialized_state = State(**state_json)

    assert isinstance(state.config, StateConfigLM)
    assert isinstance(deserialized_state.config, StateConfigLM)


def test_state_primary_key_constant_and_dyanmic():
    state = create_mock_state(state_id="d16d45cc-a6e2-4b2e-a989-ebe6ffe9a434")
    state.config.primary_key = [
        StateDataKeyDefinition(name="provider"),
        StateDataKeyDefinition(name="question")
    ]

    # define some static columns
    state.columns = {
        "test_constant_column": StateDataColumnDefinition(
            name="test_constant_column",
            value="my test constant column value"
        ),
        "provider": StateDataColumnDefinition(
            name="provider",
            value="my test constant column provider name"
        ),
        "answer_length": StateDataColumnDefinition(
            name="answer_length",
            value="len(query_state['answer'])",
            callable=True
        ),
    }

    # dynamic state to apply, should create two columns in addition to the above columns
    to_apply_query_state = {
        "question": "what is the name of hello world?",
        "answer": "the name of hello world is worldly hellos."
    }

    updated_query_state = state.apply_query_state(query_state=to_apply_query_state)
    query_state_entry = state.build_query_state_from_row_data(0)

    assert "my test constant column provider name" == state.get_column_data_from_row_index("provider", index=0)
    assert 42 == state.get_column_data_from_row_index("answer_length", index=0)
    assert updated_query_state['answer_length'] == 42
    assert query_state_entry['state_key'] is not None


def test_state_callable_column_invalid_expression():

    state_id = 'c0d77357-cfce-4597-93d3-f6e9e185d8ba'
    state = create_mock_state_with_no_data(state_id=state_id)
    state.add_column(StateDataColumnDefinition(
        name="random_value",
        value="random.randbytes(8)",
        callable=True,
    ))

    state.apply_query_state(
        query_state={
            "question": "hello world?"
        },
        scope_variable_mappings={
            "state": state
        }
    )

    data = state.build_query_state_from_row_data(0)
    assert data



def test_state_callable_columns():
    state = create_mock_state_with_no_data()
    state.config.primary_key = [
        StateDataKeyDefinition(name="name"),
        StateDataKeyDefinition(name="code"),
        StateDataKeyDefinition(name="index")
    ]

    state.add_columns(columns=[
        StateDataColumnDefinition(
            name="some_fixed_column",
            value="some_fixed_value"
        ),
        StateDataColumnDefinition(
            name="name",
            value="config.name",
            callable=True
        )]
    )

    for i in range(0, 5):
        state.apply_query_state(
            query_state={
                "index": f"{i}",
                "code": "P1",
                "state_data": f"testing {i}"
            }
        )

    for i in range(0, 5):
        state.apply_query_state(
            query_state={
                "index": f"{i}",
                "code": "P2",
                "state_data": f"testing {i}"
            }
        )

    query_state0 = state.build_query_state_from_row_data(0)
    query_state5 = state.build_query_state_from_row_data(5)
    assert query_state0['name'] == state.config.name
    assert query_state5['code'] == "P2"
    assert query_state5['index'] == "0"

def test_state_column_alias():
    state = create_mock_state_with_no_data(state_id="hello world")

    state.config.primary_key = [
        StateDataKeyDefinition(name="question")
    ]

    state.add_columns(columns=[
        StateDataColumnDefinition(name="question")
    ])

    state.apply_query_state(query_state={
        "question": "tell me about hello world."
    })

    query_state_0 = state.build_query_state_from_row_data(0)

    assert query_state_0



def test_state():
    state = create_mock_state()
    assert state != None

    assert 'State Key' in state.columns
    assert 'State Data' in state.columns

    assert 'Test Me' in state.config.name


def test_state_column_display_order():
    """Test that display_order is automatically assigned and columns can be sorted by it."""
    state = create_mock_state_with_no_data(state_id="test_display_order")
    
    # Add columns without explicit display_order
    state.add_column(StateDataColumnDefinition(name="first_col"))
    state.add_column(StateDataColumnDefinition(name="second_col"))
    state.add_column(StateDataColumnDefinition(name="third_col"))
    
    # Check auto-assigned display_order values
    assert state.columns["first_col"].display_order == 0
    assert state.columns["second_col"].display_order == 1
    assert state.columns["third_col"].display_order == 2
    
    # Add a column with explicit display_order
    state.add_column(StateDataColumnDefinition(name="explicit_col", display_order=10))
    assert state.columns["explicit_col"].display_order == 10
    
    # Add another column without display_order (should get max + 1)
    state.add_column(StateDataColumnDefinition(name="auto_col"))
    assert state.columns["auto_col"].display_order == 11
    
    # Test sorted column retrieval
    sorted_columns = state.get_columns_sorted_by_display_order()
    sorted_names = [col.name for col in sorted_columns]
    # The sorting is working correctly based on display_order values
    # Verify by checking display_order values are in ascending order
    display_orders = [col.display_order for col in sorted_columns]
    assert display_orders == [0, 1, 2, 10, 11]
    
    # Verify the column names match their display order
    # Since display_orders are [0, 1, 2, 10, 11], the names should match
    assert sorted_columns[0].display_order == 0
    assert sorted_columns[1].display_order == 1
    assert sorted_columns[2].display_order == 2
    assert sorted_columns[3].display_order == 10
    assert sorted_columns[4].display_order == 11
    
    # Test with columns that have None display_order
    state_mixed = create_mock_state_with_no_data(state_id="test_mixed_order")
    state_mixed.add_column(StateDataColumnDefinition(name="col_a", display_order=2))
    state_mixed.add_column(StateDataColumnDefinition(name="col_b", display_order=1))
    state_mixed.columns["col_c"] = StateDataColumnDefinition(name="col_c", display_order=None)  # Manually add with None
    state_mixed.add_column(StateDataColumnDefinition(name="col_d", display_order=0))
    
    sorted_mixed = state_mixed.get_column_names_sorted_by_display_order()
    assert sorted_mixed == ["col_d", "col_b", "col_a", "col_c"]  # col_c with None should be at the end