import pytest
from unittest.mock import patch, MagicMock

from auto_causal.components.query_interpreter import interpret_query
from auto_causal.models import LLMTreatmentReferenceLevel

# Basic mock data setup
MOCK_QUERY_INFO_REF_LEVEL = {
    "query_text": "What is the effect of different fertilizers (Nitro, Phos, Control) on crop_yield, using Control as the baseline?",
    "potential_treatments": ["fertilizer_type"],
    "outcome_hints": ["crop_yield"],
    "covariates_hints": ["soil_ph", "rainfall"]
}

MOCK_DATASET_ANALYSIS_REF_LEVEL = {
    "columns": ["fertilizer_type", "crop_yield", "soil_ph", "rainfall"],
    "column_categories": {
        "fertilizer_type": "categorical_multi", # Assuming a category type for multi-level
        "crop_yield": "continuous_numeric",
        "soil_ph": "continuous_numeric",
        "rainfall": "continuous_numeric"
    },
    "potential_treatments": ["fertilizer_type"],
    "potential_outcomes": ["crop_yield"],
    "value_counts": { # Added for providing unique values to the prompt
        "fertilizer_type": {
            "values": ["Nitro", "Phos", "Control"]
        }
    },
    "columns_data_preview": { # Fallback if value_counts isn't structured as expected
        "fertilizer_type": ["Nitro", "Phos", "Control", "Nitro", "Control"]
    }
    # Add other necessary fields from DatasetAnalysis model if interpret_query uses them
}

MOCK_DATASET_DESCRIPTION_REF_LEVEL = "A dataset from an agricultural experiment."

def test_interpret_query_identifies_treatment_reference_level():
    """
    Test that interpret_query correctly identifies and returns the treatment_reference_level
    when the LLM simulation provides one.
    """
    # Mock the LLM client and its structured output
    mock_llm_instance = MagicMock()
    mock_structured_llm = MagicMock()

    # This will be the mock for the call related to TREATMENT_REFERENCE_IDENTIFICATION_PROMPT_TEMPLATE
    # The other LLM calls (for T, O, C, IV, RDD, RCT) also need to be considered
    # or made to return benign defaults for this specific test.

    # Simulate LLM responses for different calls within interpret_query
    def mock_llm_call_router(*args, **kwargs):
        # The first argument to _call_llm_for_var is the llm instance,
        # The second is the prompt string
        # The third is the Pydantic model for structured output

        # args[0] is llm, args[1] is prompt, args[2] is pydantic_model
        pydantic_model_passed = args[2]

        if pydantic_model_passed == LLMTreatmentReferenceLevel:
            return LLMTreatmentReferenceLevel(reference_level="Control", reasoning="Identified from query text.")
        # Add mocks for other LLM calls if interpret_query strictly needs them to proceed
        # For example, for identifying treatment, outcome, covariates, IV, RDD, RCT:
        elif "most likely treatment variable" in args[1]: # Simplified check for treatment prompt
            return MagicMock(variable_name="fertilizer_type")
        elif "most likely outcome variable" in args[1]: # Simplified check for outcome prompt
             return MagicMock(variable_name="crop_yield")
        elif "valid covariates" in args[1]: # Simplified check for covariates prompt
             return MagicMock(covariates=["soil_ph", "rainfall"])
        elif "Instrumental Variables" in args[1]: # Check for IV prompt
            return MagicMock(instrument_variable=None)
        elif "Regression Discontinuity Design" in args[1]: # Check for RDD prompt
            return MagicMock(running_variable=None, cutoff_value=None)
        elif "Randomized Controlled Trial" in args[1]: # Check for RCT prompt
            return MagicMock(is_rct=False, reasoning="No indication of RCT.")
        return MagicMock() # Default mock for other calls

    # Patch _call_llm_for_var which is used internally by interpret_query's helpers
    with patch('auto_causal.components.query_interpreter._call_llm_for_var', side_effect=mock_llm_call_router) as mock_llm_call:
        # Patch get_llm_client to return our mock_llm_instance
        # This ensures that _call_llm_for_var uses the intended LLM mock when called from within interpret_query
        with patch('auto_causal.components.query_interpreter.get_llm_client', return_value=mock_llm_instance) as mock_get_llm:
            
            result = interpret_query(
                query_info=MOCK_QUERY_INFO_REF_LEVEL,
                dataset_analysis=MOCK_DATASET_ANALYSIS_REF_LEVEL,
                dataset_description=MOCK_DATASET_DESCRIPTION_REF_LEVEL
            )

    assert "treatment_reference_level" in result, "treatment_reference_level should be in the result"
    assert result["treatment_reference_level"] == "Control", "Incorrect treatment_reference_level identified"
    
    # Verify that the LLM was called to get the reference level
    # This requires checking the calls made to the mock_llm_call
    found_ref_level_call = False
    for call_args in mock_llm_call.call_args_list:
        # call_args is a tuple; call_args[0] contains positional args, call_args[1] has kwargs
        # The third positional argument to _call_llm_for_var is the pydantic_model
        if len(call_args[0]) >= 3 and call_args[0][2] == LLMTreatmentReferenceLevel:
            found_ref_level_call = True
            # Optionally, check the prompt content here too if needed
            # prompt_content = call_args[0][1]
            # assert "using Control as the baseline" in prompt_content
            break
    assert found_ref_level_call, "LLM call for treatment reference level was not made."

    # Basic checks for other essential variables (assuming they are mocked simply)
    assert result["treatment_variable"] == "fertilizer_type"
    assert result["outcome_variable"] == "crop_yield"
    assert result["is_rct"] is False # Based on mock