import importlib
import json
import sys
import types
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
from pydantic import ValidationError

from core.rag.models.document import Document


def _build_fake_tablestore_module():
    tablestore = types.ModuleType("tablestore")

    class _BatchGetRowRequest:
        def __init__(self):
            self.items = []

        def add(self, item):
            self.items.append(item)

    class _TableInBatchGetRowItem:
        def __init__(self, table_name, rows_to_get, columns_to_get, _unused, _ver):
            self.table_name = table_name
            self.rows_to_get = rows_to_get
            self.columns_to_get = columns_to_get

    class _Row:
        def __init__(self, primary_key, attribute_columns=None):
            self.primary_key = primary_key
            self.attribute_columns = attribute_columns or []

    class _Client:
        def __init__(self, *_args):
            self.list_table = MagicMock(return_value=[])
            self.create_table = MagicMock()
            self.list_search_index = MagicMock(return_value=[])
            self.create_search_index = MagicMock()
            self.delete_search_index = MagicMock()
            self.delete_table = MagicMock()
            self.put_row = MagicMock()
            self.delete_row = MagicMock()
            self.get_row = MagicMock(return_value=(None, None, None))
            self.batch_get_row = MagicMock()
            self.search = MagicMock()

    tablestore.OTSClient = _Client
    tablestore.BatchGetRowRequest = _BatchGetRowRequest
    tablestore.TableInBatchGetRowItem = _TableInBatchGetRowItem
    tablestore.Row = _Row
    tablestore.TableMeta = lambda name, schema: ("table_meta", name, schema)
    tablestore.TableOptions = lambda: ("table_options",)
    tablestore.CapacityUnit = lambda read, write: ("capacity", read, write)
    tablestore.ReservedThroughput = lambda cap: ("reserved", cap)
    tablestore.FieldSchema = lambda *args, **kwargs: ("field", args, kwargs)
    tablestore.VectorOptions = lambda **kwargs: ("vector_options", kwargs)
    tablestore.SearchIndexMeta = lambda field_schemas: ("search_index_meta", field_schemas)
    tablestore.SearchQuery = lambda query, **kwargs: SimpleNamespace(query=query, **kwargs)
    tablestore.TermQuery = lambda key, value: ("term_query", key, value)
    tablestore.ColumnsToGet = lambda **kwargs: ("columns_to_get", kwargs)
    tablestore.KnnVectorQuery = lambda **kwargs: SimpleNamespace(**kwargs)
    tablestore.TermsQuery = lambda key, values: ("terms_query", key, values)
    tablestore.Sort = lambda **kwargs: ("sort", kwargs)
    tablestore.ScoreSort = lambda **kwargs: ("score_sort", kwargs)
    tablestore.BoolQuery = lambda **kwargs: SimpleNamespace(**kwargs)
    tablestore.MatchQuery = lambda **kwargs: ("match_query", kwargs)

    tablestore.FieldType = SimpleNamespace(TEXT="TEXT", VECTOR="VECTOR", KEYWORD="KEYWORD")
    tablestore.AnalyzerType = SimpleNamespace(MAXWORD="MAXWORD")
    tablestore.VectorDataType = SimpleNamespace(VD_FLOAT_32="VD_FLOAT_32")
    tablestore.VectorMetricType = SimpleNamespace(VM_COSINE="VM_COSINE")
    tablestore.ColumnReturnType = SimpleNamespace(SPECIFIED="SPECIFIED", ALL_FROM_INDEX="ALL_FROM_INDEX")
    tablestore.SortOrder = SimpleNamespace(DESC="DESC")
    return tablestore


@pytest.fixture
def tablestore_module(monkeypatch: pytest.MonkeyPatch):
    fake_module = _build_fake_tablestore_module()
    monkeypatch.setitem(sys.modules, "tablestore", fake_module)

    import dify_vdb_tablestore.tablestore_vector as module

    return importlib.reload(module)


def _config(module, **overrides):
    values = {
        "access_key_id": "ak",
        "access_key_secret": "sk",
        "instance_name": "instance",
        "endpoint": "endpoint",
        "normalize_full_text_bm25_score": False,
    }
    values.update(overrides)
    return module.TableStoreConfig.model_validate(values)


@pytest.mark.parametrize(
    ("field", "value", "message"),
    [
        ("access_key_id", "", "config ACCESS_KEY_ID is required"),
        ("access_key_secret", "", "config ACCESS_KEY_SECRET is required"),
        ("instance_name", "", "config INSTANCE_NAME is required"),
        ("endpoint", "", "config ENDPOINT is required"),
    ],
)
def test_tablestore_config_validation(tablestore_module, field, value, message):
    values = _config(tablestore_module).model_dump()
    values[field] = value
    with pytest.raises(ValidationError, match=message):
        tablestore_module.TableStoreConfig.model_validate(values)


def test_init_and_basic_delegation(tablestore_module):
    vector = tablestore_module.TableStoreVector("collection_1", _config(tablestore_module))
    assert vector.get_type() == tablestore_module.VectorType.TABLESTORE
    assert vector._table_name == "collection_1"
    assert vector._index_name == "collection_1_idx"

    vector._create_collection = MagicMock()
    vector.add_texts = MagicMock()
    docs = [Document(page_content="hello", metadata={"doc_id": "d-1"})]
    vector.create(docs, [[0.1, 0.2]])
    vector._create_collection.assert_called_once_with(2)
    vector.add_texts.assert_called_once_with(documents=docs, embeddings=[[0.1, 0.2]])

    vector.create_collection([[0.1, 0.2]])
    assert vector._create_collection.call_count == 2


def test_get_by_ids_text_exists_delete_and_wrappers(tablestore_module):
    vector = tablestore_module.TableStoreVector("collection_1", _config(tablestore_module))

    # get_by_ids
    ok_item = SimpleNamespace(
        is_ok=True,
        row=SimpleNamespace(
            attribute_columns=[("metadata", json.dumps({"doc_id": "1"}), None), ("page_content", "text-1", None)]
        ),
    )
    fail_item = SimpleNamespace(is_ok=False, row=None)
    batch_resp = SimpleNamespace(get_result_by_table=lambda _table: [ok_item, fail_item])
    vector._tablestore_client.batch_get_row.return_value = batch_resp
    docs = vector.get_by_ids(["id-1"])
    assert len(docs) == 1
    assert docs[0].page_content == "text-1"

    # text_exists
    vector._tablestore_client.get_row.return_value = (None, object(), None)
    assert vector.text_exists("id-1") is True
    vector._tablestore_client.get_row.return_value = (None, None, None)
    assert vector.text_exists("id-1") is False

    # delete wrappers
    vector._delete_row = MagicMock()
    vector.delete_by_ids([])
    vector._delete_row.assert_not_called()
    vector.delete_by_ids(["id-1", "id-2"])
    assert vector._delete_row.call_count == 2

    vector._search_by_metadata = MagicMock(return_value=["id-a"])
    assert vector.get_ids_by_metadata_field("document_id", "doc-1") == ["id-a"]
    vector.delete_by_ids = MagicMock()
    vector.delete_by_metadata_field("document_id", "doc-1")
    vector.delete_by_ids.assert_called_once_with(["id-a"])

    vector._search_by_vector = MagicMock(return_value=["vec-doc"])
    vector._search_by_full_text = MagicMock(return_value=["fts-doc"])
    assert vector.search_by_vector([0.1], top_k=2, score_threshold=0.5, document_ids_filter=["d-1"]) == ["vec-doc"]
    assert vector.search_by_full_text("query", top_k=2, score_threshold=0.3, document_ids_filter=["d-1"]) == ["fts-doc"]

    vector._delete_table_if_exist = MagicMock()
    vector.delete()
    vector._delete_table_if_exist.assert_called_once()


def test_create_collection_and_table_index_lifecycle(tablestore_module, monkeypatch: pytest.MonkeyPatch):
    vector = tablestore_module.TableStoreVector("collection_1", _config(tablestore_module))
    lock = MagicMock()
    lock.__enter__.return_value = None
    lock.__exit__.return_value = None
    monkeypatch.setattr(tablestore_module.redis_client, "lock", MagicMock(return_value=lock))
    monkeypatch.setattr(tablestore_module.redis_client, "set", MagicMock())

    monkeypatch.setattr(tablestore_module.redis_client, "get", MagicMock(return_value=1))
    vector._create_table_if_not_exist = MagicMock()
    vector._create_search_index_if_not_exist = MagicMock()
    vector._create_collection(3)
    vector._create_table_if_not_exist.assert_not_called()

    monkeypatch.setattr(tablestore_module.redis_client, "get", MagicMock(return_value=None))
    vector._create_collection(3)
    vector._create_table_if_not_exist.assert_called_once()
    vector._create_search_index_if_not_exist.assert_called_once_with(3)
    tablestore_module.redis_client.set.assert_called_once()

    vector = tablestore_module.TableStoreVector("collection_2", _config(tablestore_module))
    vector._tablestore_client.list_table.return_value = ["collection_2"]
    assert vector._create_table_if_not_exist() is None
    vector._tablestore_client.list_table.return_value = []
    vector._create_table_if_not_exist()
    vector._tablestore_client.create_table.assert_called_once()

    vector._tablestore_client.list_search_index.return_value = [("collection_2", "collection_2_idx")]
    assert vector._create_search_index_if_not_exist(3) is None
    vector._tablestore_client.list_search_index.return_value = []
    vector._create_search_index_if_not_exist(3)
    vector._tablestore_client.create_search_index.assert_called_once()

    vector._tablestore_client.list_search_index.return_value = [("collection_2", "idx_a"), ("collection_2", "idx_b")]
    vector._delete_table_if_exist()
    assert vector._tablestore_client.delete_search_index.call_count == 2
    vector._tablestore_client.delete_table.assert_called_once_with("collection_2")

    vector._delete_search_index()
    vector._tablestore_client.delete_search_index.assert_called_with("collection_2", "collection_2_idx")


def test_write_row_and_search_helpers(tablestore_module):
    vector = tablestore_module.TableStoreVector("collection_1", _config(tablestore_module))

    vector._write_row(
        "id-1",
        {
            "page_content": "hello",
            "vector": [0.1, 0.2],
            "metadata": {"doc_id": "d-1", "document_id": "doc-1"},
        },
    )
    put_row_call = vector._tablestore_client.put_row.call_args
    assert put_row_call.args[0] == "collection_1"
    attrs = put_row_call.args[1].attribute_columns
    assert any(item[0] == "metadata_tags" for item in attrs)

    vector._delete_row("id-1")
    vector._tablestore_client.delete_row.assert_called_once()

    # metadata search pagination
    first_page = SimpleNamespace(rows=[[(("id", "row-1"),)]], next_token=b"next")
    second_page = SimpleNamespace(rows=[[(("id", "row-2"),)]], next_token=b"")
    vector._tablestore_client.search.side_effect = [first_page, second_page]
    ids = vector._search_by_metadata("document_id", "doc-1")
    assert ids == ["row-1", "row-2"]
    vector._tablestore_client.search.side_effect = None

    # vector search
    hit1 = SimpleNamespace(
        score=0.9,
        row=(
            None,
            [("page_content", "doc-a"), ("metadata", json.dumps({"doc_id": "1"})), ("vector", json.dumps([0.1]))],
        ),
    )
    hit2 = SimpleNamespace(
        score=0.2,
        row=(
            None,
            [("page_content", "doc-b"), ("metadata", json.dumps({"doc_id": "2"})), ("vector", json.dumps([0.2]))],
        ),
    )
    vector._tablestore_client.search.return_value = SimpleNamespace(search_hits=[hit1, hit2])
    docs = vector._search_by_vector([0.1], document_ids_filter=["document_id=doc-1"], top_k=2, score_threshold=0.5)
    assert len(docs) == 1
    assert docs[0].metadata["score"] == pytest.approx(0.9)

    assert tablestore_module.TableStoreVector._normalize_score_exp_decay(0) == pytest.approx(0.0)
    assert tablestore_module.TableStoreVector._normalize_score_exp_decay(100) <= 1.0

    # full text search with and without normalized score filter
    vector._normalize_full_text_bm25_score = True
    hit3 = SimpleNamespace(
        score=10.0, row=(None, [("page_content", "doc-c"), ("metadata", json.dumps({"doc_id": "3"}))])
    )
    hit4 = SimpleNamespace(
        score=0.1, row=(None, [("page_content", "doc-d"), ("metadata", json.dumps({"doc_id": "4"}))])
    )
    vector._tablestore_client.search.return_value = SimpleNamespace(search_hits=[hit3, hit4])
    docs = vector._search_by_full_text("query", document_ids_filter=["document_id=doc-1"], top_k=2, score_threshold=0.2)
    assert len(docs) == 1
    assert "score" in docs[0].metadata

    vector._normalize_full_text_bm25_score = False
    vector._tablestore_client.search.return_value = SimpleNamespace(search_hits=[hit3])
    docs = vector._search_by_full_text("query", document_ids_filter=None, top_k=2, score_threshold=0.0)
    assert len(docs) == 1
    assert "score" not in docs[0].metadata


def test_tablestore_factory_uses_existing_or_generated_collection(tablestore_module, monkeypatch: pytest.MonkeyPatch):
    factory = tablestore_module.TableStoreVectorFactory()
    dataset_with_index = SimpleNamespace(
        id="dataset-1",
        index_struct_dict={"vector_store": {"class_prefix": "EXISTING_COLLECTION"}},
        index_struct=None,
    )
    dataset_without_index = SimpleNamespace(id="dataset-2", index_struct_dict=None, index_struct=None)

    monkeypatch.setattr(tablestore_module.Dataset, "gen_collection_name_by_id", lambda _id: "AUTO_COLLECTION")
    monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_ENDPOINT", "endpoint")
    monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_INSTANCE_NAME", "instance")
    monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_ACCESS_KEY_ID", "ak")
    monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_ACCESS_KEY_SECRET", "sk")
    monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_NORMALIZE_FULLTEXT_BM25_SCORE", True)

    with patch.object(tablestore_module, "TableStoreVector", return_value="vector") as vector_cls:
        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())

    assert result_1 == "vector"
    assert result_2 == "vector"
    assert vector_cls.call_args_list[0].kwargs["collection_name"] == "EXISTING_COLLECTION"
    assert vector_cls.call_args_list[1].kwargs["collection_name"] == "AUTO_COLLECTION"
    assert dataset_without_index.index_struct is not None
