import pytest

from monarch_py.datamodels.category_enums import (
    AssociationCategory,
    AssociationPredicate,
)
from monarch_py.datamodels.model import Node
from monarch_py.implementations.solr.solr_query_utils import (
    build_association_counts_query,
    build_association_query,
    build_autocomplete_query,
    build_histopheno_query,
    build_mapping_query,
    build_search_query,
    obsolete_unboost,
    entity_predicate_boost,
    category_boost,
    blank_search_boost,
    build_association_count_suffixes,
    build_grid_column_query,
    build_grid_row_query,
    build_multi_category_column_query,
    build_multi_category_row_query,
)
from monarch_py.datamodels.grid_configs import get_grid_config
from monarch_py.datamodels.grid_groupings import RowGroupingConfig, GroupingType
from monarch_py.utils.utils import compare_dicts, dict_diff


@pytest.mark.parametrize(
    "direct, facet_fields, facet_queries",
    [
        (True, None, None),
        (False, None, None),
        (None, ["category", "predicate"], [AssociationCategory.DISEASE_TO_PHENOTYPIC_FEATURE_ASSOCIATION.value]),
        (None, ["category", "predicate"], [AssociationCategory.DISEASE_TO_PHENOTYPIC_FEATURE_ASSOCIATION.value]),
    ],
)
def test_build_association_query(
    direct,
    facet_fields,
    facet_queries,
    association_query_params,
    association_query_direct,
    association_query_indirect,
):
    query = build_association_query(
        **association_query_params,
        direct=direct,
        facet_fields=facet_fields,
        facet_queries=facet_queries,
    ).model_dump()
    expected = association_query_direct if direct else association_query_indirect
    if facet_fields:
        expected["facet_fields"] = facet_fields
    if facet_queries:
        expected["facet_queries"] = facet_queries
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_build_association_multiple_categories():
    query = build_association_query(
        category=[
            AssociationCategory.CAUSAL_GENE_TO_DISEASE_ASSOCIATION.value,
            AssociationCategory.DISEASE_TO_PHENOTYPIC_FEATURE_ASSOCIATION.value,
        ]
    )
    assert len(query.filter_queries) > 0, "filter_queries is empty"
    category_filter = [fq for fq in query.filter_queries if fq.startswith("category:")][0]
    assert (
        category_filter
        == "category:biolink\\:CausalGeneToDiseaseAssociation OR category:biolink\\:DiseaseToPhenotypicFeatureAssociation"
    ), "multiple category filter is not as expected"


def test_build_association_multiple_predicates():
    query = build_association_query(
        predicate=[AssociationPredicate.HAS_PHENOTYPE.value, AssociationPredicate.EXPRESSED_IN.value]
    )
    assert len(query.filter_queries) > 0, "filter_queries is empty"
    predicate_filter = [fq for fq in query.filter_queries if fq.startswith("predicate:")][0]
    assert predicate_filter == "predicate:biolink\\:has_phenotype OR predicate:biolink\\:expressed_in", (
        "multiple predicate filter is not as expected"
    )


def test_build_association_multiple_entites():
    query = build_association_query(entity=["MONDO:0020121", "HP:0000006"])
    assert len(query.filter_queries) > 0, "filter_queries is empty"
    entity_filter = [fq for fq in query.filter_queries if fq.startswith("subject:")][0]
    assert (
        entity_filter
        == 'subject:"MONDO:0020121" OR subject_closure:"MONDO:0020121" OR object:"MONDO:0020121" OR object_closure:"MONDO:0020121" OR disease_context_qualifier:"MONDO:0020121" OR disease_context_qualifier_closure:"MONDO:0020121" OR subject:"HP:0000006" OR subject_closure:"HP:0000006" OR object:"HP:0000006" OR object_closure:"HP:0000006" OR disease_context_qualifier:"HP:0000006" OR disease_context_qualifier_closure:"HP:0000006"'
    )


def test_build_association_multiple_subjects():
    query = build_association_query(subject=["MONDO:0020121", "MONDO:0007915"])
    assert len(query.filter_queries) > 0, "filter_queries is empty"
    subject_filter = [fq for fq in query.filter_queries if fq.startswith("subject:")][0]
    assert (
        subject_filter
        == 'subject:"MONDO:0020121" OR subject_closure:"MONDO:0020121" OR subject:"MONDO:0007915" OR subject_closure:"MONDO:0007915"'
    ), "multiple subject filter is not as expected"


def test_build_association_multiple_objects():
    query = build_association_query(object=["HP:0000006", "HP:0000007"])
    assert len(query.filter_queries) > 0, "filter_queries is empty"
    object_filter = [fq for fq in query.filter_queries if fq.startswith("object:")][0]
    print(object_filter)
    assert (
        object_filter
        == 'object:"HP:0000006" OR object_closure:"HP:0000006" OR object:"HP:0000007" OR object_closure:"HP:0000007"'
    ), "multiple object filter is not as expected"


def test_build_association_counts_query(association_counts_query, node):
    query = build_association_counts_query(entities=[Node(**node).id]).model_dump()
    expected = association_counts_query
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_build_association_counts_query_with_orthologs():
    """Test that multi-entity queries include ortho_subject/ortho_object facet queries."""
    entities = ["HGNC:1100", "MGI:88276", "ZFIN:ZDB-GENE-040426-2889"]
    query = build_association_counts_query(entities=entities)

    # Should have 6 groups of facet queries: direct_subject, direct_object,
    # closure_subject, closure_object, ortho_subject, ortho_object
    from monarch_py.utils.association_type_utils import AssociationTypeMappings

    n_mappings = len(AssociationTypeMappings.get_mappings())
    assert len(query.facet_queries) == n_mappings * 6

    # Verify ortholog facet queries reference all entities
    ortho_queries = query.facet_queries[n_mappings * 4 :]
    for fq in ortho_queries:
        for ent in entities:
            assert ent in fq, f"Expected entity {ent} in ortholog facet query: {fq}"


def test_build_histopheno_query(histopheno_query):
    query = build_histopheno_query("MONDO:0020121").model_dump()
    expected = histopheno_query
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_build_search_query(search_query):
    query = build_search_query(q="fanconi").model_dump()
    expected = search_query
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_build_autocomplete_query(autocomplete_query):
    query = build_autocomplete_query(q="fanc").model_dump()
    expected = autocomplete_query
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_build_mappings_query(mapping_query):
    query = build_mapping_query(entity_id=["MONDO:0020121"]).model_dump()
    expected = mapping_query
    assert compare_dicts(query, expected), f"Query is not as expected. Difference: {dict_diff(query, expected)}"


def test_obsolete_unboost():
    boost = obsolete_unboost()
    assert "deprecated" in boost


def test_entity_predicate_boost() -> str:
    boost = entity_predicate_boost([AssociationPredicate.HAS_PHENOTYPE], 2.0)
    assert "has_phenotype_count" in boost
    assert "2.0" in boost


def test_category_boost():
    boost = category_boost("biolink:PhenotypicFeature", 99.0)
    assert "biolink:PhenotypicFeature" in boost
    assert "99.0" in boost


def test_category_boost_with_taxon():
    boost = category_boost("biolink:Gene", 100.0, taxon="NCBITaxon:9606")
    assert "biolink:Gene" in boost
    assert "100.0" in boost
    assert "NCBITaxon:9606" in boost


def test_blank_search_boost():
    boost = blank_search_boost()
    assert boost.startswith("product(")
    assert boost.endswith(")")
    assert 'if(termfreq(id,"MONDO:0007523"),12,1)' in boost


def test_build_association_query_with_primary_knowledge_source():
    """Test that primary_knowledge_source param adds the correct filter query."""
    query = build_association_query(
        primary_knowledge_source=["infores:hpo-annotations"],
    )
    pks_filters = [fq for fq in query.filter_queries if "primary_knowledge_source" in fq]
    assert len(pks_filters) == 1, "Expected one primary_knowledge_source filter"
    assert "infores" in pks_filters[0]


def test_build_association_query_with_q_parameter():
    """Test that q param enables edismax search with highlighting."""
    query = build_association_query(q="BRCA1")
    assert query.q == "BRCA1"
    assert query.def_type == "edismax"
    assert query.hl is True


# =====================================================================
# Tests for AssociationCountSuffixes
# =====================================================================


@pytest.mark.parametrize(
    "entities,expect_orthologs",
    [
        (["HGNC:1100"], False),
        (["HGNC:1100", "MGI:88276"], True),
    ],
    ids=["single-entity", "multi-entity"],
)
def test_suffixes_ortholog_presence(entities, expect_orthologs):
    suffixes = build_association_count_suffixes(entities)
    all_s = suffixes.all_suffixes
    assert "direct_subject" in all_s
    assert "direct_object" in all_s
    assert "closure_subject" in all_s
    assert "closure_object" in all_s
    assert ("orthologs_subject" in all_s) == expect_orthologs
    assert ("orthologs_object" in all_s) == expect_orthologs
    if expect_orthologs:
        assert len(all_s) == 6


def test_suffix_strings_contain_entity():
    suffixes = build_association_count_suffixes(["HGNC:1100"])
    assert 'subject:"HGNC:1100"' in suffixes.direct_subject
    assert 'object:"HGNC:1100"' in suffixes.direct_object
    assert 'subject:"HGNC:1100"' in suffixes.closure_subject
    assert 'object:"HGNC:1100"' in suffixes.closure_object


def test_ortho_suffixes_contain_all_entities():
    entities = ["HGNC:1100", "MGI:88276", "ZFIN:ZDB"]
    suffixes = build_association_count_suffixes(entities)
    for ent in entities:
        assert ent in suffixes.ortho_subject
        assert ent in suffixes.ortho_object


# =====================================================================
# Tests for build_grid_column_query
# =====================================================================


@pytest.mark.parametrize(
    "direct_only,expected_field",
    [(True, 'object:"MONDO:0007078"'), (False, 'object_closure:"MONDO:0007078"')],
    ids=["direct", "indirect"],
)
def test_column_query_direct_vs_indirect(direct_only, expected_field):
    config = get_grid_config("case-phenotype")
    params = build_grid_column_query(
        context_id="MONDO:0007078",
        config=config,
        direct_only=direct_only,
    )
    fq_str = str(params.get("fq", ""))
    assert expected_field in fq_str


@pytest.mark.parametrize(
    "filter_empty,expect_join",
    [(True, True), (False, False)],
    ids=["with-join", "without-join"],
)
def test_column_query_join_filtering(filter_empty, expect_join):
    config = get_grid_config("case-phenotype")
    params = build_grid_column_query(
        context_id="MONDO:0007078",
        config=config,
        direct_only=True,
        filter_empty_columns=filter_empty,
    )
    assert ("{!join" in params["q"]) == expect_join


# =====================================================================
# Tests for build_grid_row_query
# =====================================================================


def test_row_query_structure():
    config = get_grid_config("case-phenotype")
    grouping = RowGroupingConfig(
        grouping_type=GroupingType.CLOSURE_ROOTS,
        bin_ids=["BIN:001"],
        bin_labels={"BIN:001": "Bin 1"},
    )
    params = build_grid_row_query(
        context_id="MONDO:0007078",
        config=config,
        grouping=grouping,
        direct_only=True,
    )
    assert "{!join" in params["q"]
    assert 'category:"biolink:CaseToPhenotypicFeatureAssociation"' in params["fq"]
    assert params["facet"] == "true"
    assert 'object_closure:"BIN:001"' in params["facet.query"]


def test_row_query_indirect():
    config = get_grid_config("case-phenotype")
    grouping = RowGroupingConfig(
        grouping_type=GroupingType.CLOSURE_ROOTS,
        bin_ids=["BIN:001"],
        bin_labels={"BIN:001": "Bin 1"},
    )
    params = build_grid_row_query(
        context_id="MONDO:0007078",
        config=config,
        grouping=grouping,
        direct_only=False,
    )
    assert "object_closure" in params["q"]


# =====================================================================
# Tests for build_multi_category_column_query
# =====================================================================


def test_multi_category_column_query_multiple_categories():
    params = build_multi_category_column_query(
        context_id="HGNC:4851",
        column_assoc_categories=[
            "biolink:CausalGeneToDiseaseAssociation",
            "biolink:CorrelatedGeneToDiseaseAssociation",
        ],
        context_field="subject",
        context_closure_field="subject_closure",
        column_field="object",
        direct_only=True,
    )
    fq_str = str(params.get("fq", ""))
    assert "CausalGeneToDiseaseAssociation" in fq_str
    assert "CorrelatedGeneToDiseaseAssociation" in fq_str
    assert " OR " in fq_str


@pytest.mark.parametrize(
    "filter_empty,expect_join",
    [(True, True), (False, False)],
    ids=["with-join", "without-join"],
)
def test_multi_category_column_query_join_filtering(filter_empty, expect_join):
    extra_kwargs = {}
    if filter_empty:
        extra_kwargs["row_assoc_categories"] = ["biolink:DiseaseToPhenotypicFeatureAssociation"]
        extra_kwargs["row_context_field"] = "subject"
    params = build_multi_category_column_query(
        context_id="HGNC:4851",
        column_assoc_categories=["biolink:CausalGeneToDiseaseAssociation"],
        context_field="subject",
        context_closure_field="subject_closure",
        column_field="object",
        direct_only=True,
        filter_empty_columns=filter_empty,
        **extra_kwargs,
    )
    assert ("{!join" in params["q"]) == expect_join


def test_multi_category_column_query_includes_source_fields():
    params = build_multi_category_column_query(
        context_id="HGNC:4851",
        column_assoc_categories=["biolink:CausalGeneToDiseaseAssociation"],
        context_field="subject",
        context_closure_field="subject_closure",
        column_field="object",
        direct_only=True,
    )
    fl = params["fl"]
    assert "category" in fl
    assert "predicate" in fl
    assert "publications" in fl
    assert "has_evidence" in fl


def test_build_multi_category_row_query_multiple_row_categories():
    """Query builder should accept multiple row categories and build OR filter."""
    grouping = RowGroupingConfig(
        grouping_type=GroupingType.CLOSURE_ROOTS,
        bin_ids=["HP:0000001"],
        bin_labels={"HP:0000001": "All"},
    )

    query_params = build_multi_category_row_query(
        context_id="MONDO:0005148",
        column_assoc_categories=["biolink:DiseaseToPhenotypicFeatureAssociation"],
        row_assoc_categories=[
            "biolink:DiseaseToPhenotypicFeatureAssociation",
            "biolink:CaseToPhenotypicFeatureAssociation",
        ],
        context_field="object",
        context_closure_field="object_closure",
        column_field="subject",
        row_context_field="subject",
        row_entity_field="object",
        grouping=grouping,
        direct_only=True,
    )

    # Should build OR filter for multiple categories
    fq = query_params.get("fq", "")
    assert "biolink:DiseaseToPhenotypicFeatureAssociation" in fq
    assert "biolink:CaseToPhenotypicFeatureAssociation" in fq
    assert " OR " in fq


def test_build_multi_category_row_query_single_row_category_in_list():
    """Single row category passed as list should work."""
    grouping = RowGroupingConfig(
        grouping_type=GroupingType.CLOSURE_ROOTS,
        bin_ids=["HP:0000001"],
        bin_labels={"HP:0000001": "All"},
    )

    query_params = build_multi_category_row_query(
        context_id="MONDO:0005148",
        column_assoc_categories=["biolink:DiseaseToPhenotypicFeatureAssociation"],
        row_assoc_categories=["biolink:DiseaseToPhenotypicFeatureAssociation"],
        context_field="object",
        context_closure_field="object_closure",
        column_field="subject",
        row_context_field="subject",
        row_entity_field="object",
        grouping=grouping,
        direct_only=True,
    )

    fq = query_params.get("fq", "")
    assert "biolink:DiseaseToPhenotypicFeatureAssociation" in fq
    # Single category should still be wrapped in parens (from OR join)
    assert fq == '(category:"biolink:DiseaseToPhenotypicFeatureAssociation")'
