#! /usr/bin/env python3
# -*- coding: utf-8 -*-

"""
LanceDB 检索器实现

薄封装：核心价值是 embedding 集成 + Document 统一抽象。
高级操作（FTS、标量索引、computed columns）请直接使用 retriever.table 原生 API。
"""

import json
from typing import Dict, List, Optional, Type, Union, Literal

import pyarrow as pa
from loguru import logger

from ..embedding.base import BaseEmbedding
from .document import Document, SearchResult, Modality, _content_hash


DistanceMetric = Literal["cosine", "l2"]

# Python 类型 → Arrow 类型映射
_PY_TO_ARROW = {
    str: pa.string(),
    int: pa.int64(),
    float: pa.float64(),
    bool: pa.bool_(),
}

# Arrow 类型 → Python 类型映射（用于从已有 schema 自动推断）
_ARROW_TO_PY = {
    pa.string(): str,
    pa.int64(): int,
    pa.int32(): int,
    pa.float64(): float,
    pa.float32(): float,
    pa.bool_(): bool,
}

# schema 中的保留字段名
_RESERVED_FIELDS = {"id", "content", "modality", "metadata", "vector"}


class LanceRetriever:
    """
    基于 LanceDB 的检索器

    薄封装，核心提供：
    - embedding 自动调用（search/upsert 时自动 embed）
    - Document/SearchResult 统一抽象（兼容 Chroma/Milvus retriever）
    - extra_fields 将 metadata 字段提升为独立列，支持原生 SQL 过滤
    - 通过 .table 属性暴露底层 LanceDB Table，支持全部原生 API

    Examples:
        # 基础用法
        retriever = LanceRetriever(embedding, "./db", "docs")
        retriever.search("query text", top_k=5)

        # 带过滤字段
        retriever = LanceRetriever(embedding, "./db", "docs",
            extra_fields={"category": str, "price": float})
        doc = Document.text("hello", category="tech", price=29.9)
        retriever.upsert(doc)
        retriever.search("query", where="category = 'tech' AND price > 20")

        # 原生 API（FTS、标量索引等）
        retriever.table.create_scalar_index("category", index_type="BITMAP")
        retriever.table.create_fts_index("content")
        retriever.table.search("keyword", query_type="fts").limit(10).to_pandas()
    """

    def __init__(
        self,
        embedding: BaseEmbedding,
        persist_dir: str,
        table_name: str = "default",
        distance_metric: DistanceMetric = "cosine",
        extra_fields: Optional[Dict[str, Type]] = None,
    ):
        """
        Args:
            embedding: Embedding 实例（TextEmbedding 或 MultiModalEmbedding）
            persist_dir: 持久化目录
            table_name: 表名称
            distance_metric: 距离度量方式 (cosine/l2)
            extra_fields: metadata 中需要提升为独立列的字段及类型，如 {"category": str, "price": float}。
                          打开已有表时会自动从 schema 推断，无需重复指定。
        """
        import lancedb

        self.embedding = embedding
        self.persist_dir = persist_dir
        self.table_name = table_name
        self.distance_metric = distance_metric
        self._extra_fields: Dict[str, Type] = extra_fields or {}

        self.db = lancedb.connect(persist_dir)
        self._table = None
        self._dim = None

        # 尝试打开已有表
        if table_name in self.db.table_names():
            self._table = self.db.open_table(table_name)
            # 从已有 schema 推断 extra_fields（合并用户指定的）
            detected = self._detect_extra_fields()
            detected.update(self._extra_fields)
            self._extra_fields = detected
            logger.info(f"Table '{table_name}' opened, {self.count()} documents"
                        + (f", extra_fields={list(self._extra_fields)}" if self._extra_fields else ""))
        else:
            logger.info(f"Table '{table_name}' will be created on first insert")

    @property
    def table(self):
        """暴露底层 LanceDB Table，支持原生 API 操作

        Examples:
            # 标量索引
            retriever.table.create_scalar_index("category", index_type="BITMAP")
            # 全文搜索
            retriever.table.create_fts_index("content")
            retriever.table.search("keyword", query_type="fts").limit(10).to_pandas()
            # 纯 SQL 过滤（不走向量）
            retriever.table.search().where("price > 50").select(["content", "price"]).to_pandas()
        """
        return self._table

    def _detect_extra_fields(self) -> Dict[str, Type]:
        """从已有表 schema 自动推断 extra_fields"""
        if self._table is None:
            return {}
        detected = {}
        for field in self._table.schema:
            if field.name not in _RESERVED_FIELDS:
                py_type = _ARROW_TO_PY.get(field.type)
                if py_type is not None:
                    detected[field.name] = py_type
        return detected

    def _get_or_create_table(self, dim: int):
        """获取或创建表"""
        if self._table is not None:
            return self._table

        self._dim = dim
        fields = [
            pa.field("id", pa.string()),
            pa.field("content", pa.string()),
            pa.field("modality", pa.string()),
            pa.field("metadata", pa.string()),  # 未提升的 metadata 字段
        ]
        # 添加 extra_fields 为独立列
        for name, py_type in self._extra_fields.items():
            arrow_type = _PY_TO_ARROW.get(py_type)
            if arrow_type is None:
                raise ValueError(f"Unsupported field type: {name}={py_type}, supported: {list(_PY_TO_ARROW)}")
            fields.append(pa.field(name, arrow_type))

        fields.append(pa.field("vector", pa.list_(pa.float32(), dim)))

        schema = pa.schema(fields)
        self._table = self.db.create_table(self.table_name, schema=schema)
        logger.info(f"Table '{self.table_name}' created with dim={dim}"
                    + (f", extra_fields={list(self._extra_fields)}" if self._extra_fields else ""))
        return self._table

    def _embed_documents(self, documents: List[Document]) -> List[List[float]]:
        """对文档进行向量化"""
        if not documents:
            return []

        has_image = any(doc.is_image for doc in documents)
        if has_image and not self.embedding.supports_image:
            raise ValueError("Embedding 不支持图片，但文档中包含图片。请使用 MultiModalEmbedding。")

        if has_image:
            embeddings = []
            for doc in documents:
                input_type = "image" if doc.modality == "image" else "text"
                vec = self.embedding.embed([doc.content], input_type=input_type)[0]
                embeddings.append(vec)
            return embeddings
        else:
            contents = [doc.content for doc in documents]
            return self.embedding.embed(contents)

    def _embed_query(self, query: str, query_type: Modality = "text") -> List[float]:
        """对查询进行向量化"""
        if query_type == "image" and not self.embedding.supports_image:
            raise ValueError("Embedding 不支持图片查询")

        if self.embedding.supports_image:
            input_type = "image" if query_type == "image" else "text"
            return self.embedding.embed([query], input_type=input_type)[0]
        else:
            return self.embedding.embed([query])[0]

    def _embed_queries(self, queries: List[str], query_type: Modality = "text") -> List[List[float]]:
        """对多个查询进行批量向量化"""
        if not queries:
            return []

        if query_type == "image" and not self.embedding.supports_image:
            raise ValueError("Embedding 不支持图片查询")

        if self.embedding.supports_image:
            input_type = "image" if query_type == "image" else "text"
            return self.embedding.embed(queries, input_type=input_type)
        else:
            return self.embedding.embed(queries)

    def _docs_to_records(self, documents: List[Document], embeddings: List[List[float]]) -> List[dict]:
        """将文档和向量转为 LanceDB 记录，extra_fields 提升为独立列"""
        extra_names = set(self._extra_fields)
        records = []
        for doc, vec in zip(documents, embeddings):
            record = {
                "id": doc.id,
                "content": doc.content,
                "modality": doc.modality,
                "vector": vec,
            }
            # 分离 extra_fields 和剩余 metadata
            remaining = {}
            for k, v in doc.metadata.items():
                if k in extra_names:
                    record[k] = v
                else:
                    remaining[k] = v
            record["metadata"] = json.dumps(remaining, ensure_ascii=False) if remaining else ""
            records.append(record)
        return records

    def _distance_to_score(self, distance: float) -> float:
        """距离转相似度"""
        if self.distance_metric == "cosine":
            return 1 - distance
        else:
            return -distance

    # ========== 索引操作 ==========

    def add(
        self,
        documents: Union[Document, List[Document]],
        skip_existing: bool = False,
    ) -> List[str]:
        """
        添加文档

        Args:
            documents: 单个文档或文档列表
            skip_existing: 是否跳过已存在的文档

        Returns:
            添加的文档 ID 列表
        """
        if isinstance(documents, Document):
            documents = [documents]
        if not documents:
            return []

        if skip_existing:
            existing_ids = self._get_existing_ids([doc.id for doc in documents])
            documents = [doc for doc in documents if doc.id not in existing_ids]
            if not documents:
                return []

        embeddings = self._embed_documents(documents)
        records = self._docs_to_records(documents, embeddings)
        table = self._get_or_create_table(len(embeddings[0]))
        table.add(records)

        ids = [doc.id for doc in documents]
        logger.debug(f"Added {len(documents)} documents")
        return ids

    def upsert(
        self,
        documents: Union[Document, List[Document]],
        skip_existing: bool = False,
    ) -> List[str]:
        """
        添加或更新文档

        Args:
            documents: 单个文档或文档列表
            skip_existing: 是否跳过已存在的文档

        Returns:
            upsert 的文档 ID 列表
        """
        if isinstance(documents, Document):
            documents = [documents]
        if not documents:
            return []

        if skip_existing:
            existing_ids = self._get_existing_ids([doc.id for doc in documents])
            documents = [doc for doc in documents if doc.id not in existing_ids]
            if not documents:
                return []

        embeddings = self._embed_documents(documents)
        records = self._docs_to_records(documents, embeddings)
        table = self._get_or_create_table(len(embeddings[0]))

        table.merge_insert("id").when_matched_update_all().when_not_matched_insert_all().execute(records)

        ids = [doc.id for doc in documents]
        logger.debug(f"Upserted {len(documents)} documents")
        return ids

    def delete(self, ids: Union[str, List[str]]) -> None:
        """删除文档"""
        if self._table is None:
            return

        if isinstance(ids, str):
            ids = [ids]

        id_list = ", ".join(f"'{id_}'" for id_ in ids)
        self._table.delete(f"id IN ({id_list})")
        logger.debug(f"Deleted {len(ids)} documents")

    def delete_by_content(self, contents: Union[str, List[str]]) -> None:
        """根据内容删除文档"""
        if isinstance(contents, str):
            contents = [contents]

        ids = [_content_hash(content) for content in contents]
        self.delete(ids)

    # ========== 检索操作 ==========

    def search(
        self,
        query: str,
        top_k: int = 5,
        query_type: Modality = "text",
        where: Optional[str] = None,
        prefilter: bool = True,
    ) -> List[SearchResult]:
        """
        检索相似文档

        Args:
            query: 查询内容（文本或图片路径/URL）
            top_k: 返回数量
            query_type: 查询类型 "text" / "image"
            where: SQL 过滤条件，如 "category = 'tech' AND price > 20"
            prefilter: True=先过滤再搜索（默认），False=先搜索再过滤
        """
        query_embedding = self._embed_query(query, query_type)
        return self.search_by_vector(query_embedding, top_k=top_k, where=where, prefilter=prefilter)

    def search_batch(
        self,
        queries: List[str],
        top_k: int = 5,
        query_type: Modality = "text",
        where: Optional[str] = None,
        prefilter: bool = True,
    ) -> List[List[SearchResult]]:
        """批量检索相似文档"""
        if not queries:
            return []

        query_embeddings = self._embed_queries(queries, query_type)
        return self.search_by_vectors(query_embeddings, top_k=top_k, where=where, prefilter=prefilter)

    def search_by_vector(
        self,
        vector: List[float],
        top_k: int = 5,
        where: Optional[str] = None,
        prefilter: bool = True,
    ) -> List[SearchResult]:
        """
        直接使用向量检索

        Args:
            vector: 查询向量
            top_k: 返回数量
            where: SQL 过滤条件
            prefilter: True=先过滤再搜索，False=先搜索再过滤
        """
        if self._table is None:
            return []

        q = self._table.search(vector, vector_column_name="vector").metric(self.distance_metric).limit(top_k)
        if where:
            q = q.where(where, prefilter=prefilter)

        results = q.to_list()
        return self._parse_results(results)

    def search_by_vectors(
        self,
        vectors: List[List[float]],
        top_k: int = 5,
        where: Optional[str] = None,
        prefilter: bool = True,
    ) -> List[List[SearchResult]]:
        """批量使用向量检索"""
        if not vectors:
            return []

        return [self.search_by_vector(vec, top_k=top_k, where=where, prefilter=prefilter) for vec in vectors]

    def _parse_results(self, results: list) -> List[SearchResult]:
        """解析 LanceDB 查询结果，将 extra_fields 合并回 metadata"""
        extra_names = set(self._extra_fields)
        search_results = []
        for row in results:
            # 从 JSON 列恢复剩余 metadata
            raw_meta = row.get("metadata", "")
            metadata = json.loads(raw_meta) if raw_meta else {}
            # 合并 extra_fields 回 metadata
            for name in extra_names:
                if name in row:
                    metadata[name] = row[name]

            distance = row.get("_distance", 0)
            score = self._distance_to_score(distance)

            search_results.append(SearchResult(
                id=row["id"],
                content=row["content"],
                score=score,
                modality=row.get("modality", "text"),
                metadata=metadata,
            ))
        return search_results

    # ========== 管理操作 ==========

    def get(
        self,
        ids: Optional[Union[str, List[str]]] = None,
        where: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> List[Document]:
        """
        获取文档

        Args:
            ids: 文档 ID 或 ID 列表
            where: SQL 过滤条件
            limit: 返回数量限制
        """
        if self._table is None:
            return []

        table = self._table

        conditions = []
        if ids is not None:
            if isinstance(ids, str):
                ids = [ids]
            id_list = ", ".join(f"'{id_}'" for id_ in ids)
            conditions.append(f"id IN ({id_list})")
        if where:
            conditions.append(where)

        if conditions:
            combined = " AND ".join(conditions)
            results = table.search().where(combined, prefilter=True).limit(limit or 1000000).to_list()
        else:
            results = table.search().limit(limit or 1000000).to_list()

        extra_names = set(self._extra_fields)
        documents = []
        for row in results:
            raw_meta = row.get("metadata", "")
            metadata = json.loads(raw_meta) if raw_meta else {}
            for name in extra_names:
                if name in row:
                    metadata[name] = row[name]
            documents.append(Document(
                id=row["id"],
                content=row["content"],
                modality=row.get("modality", "text"),
                metadata=metadata,
            ))
        return documents

    def count(self) -> int:
        """返回文档数量"""
        if self._table is None:
            return 0
        return self._table.count_rows()

    def clear(self) -> None:
        """清空表"""
        logger.info(f"Clearing table: {self.table_name}")
        if self.table_name in self.db.table_names():
            self.db.drop_table(self.table_name)
        self._table = None
        logger.info(f"Table '{self.table_name}' cleared")

    def get_all_ids(self) -> List[str]:
        """获取所有文档 ID"""
        if self._table is None:
            return []
        results = self._table.to_pandas()["id"].tolist()
        return results

    def get_all_embeddings(self) -> tuple[list[list[float]], list[str], list[str]]:
        """获取全量向量

        Returns:
            (embeddings, ids, documents) 三元组
        """
        if self._table is None:
            return [], [], []
        df = self._table.to_pandas()
        return df["vector"].tolist(), df["id"].tolist(), df["content"].tolist()

    # ========== 便利方法 ==========

    def upsert_batch(
        self,
        documents: List[Document],
        batch_size: int = 32,
        skip_existing: bool = False,
        show_progress: bool = True,
    ) -> int:
        """
        批量插入文档（带进度条和增量更新支持）

        Args:
            documents: 文档列表
            batch_size: 批处理大小
            skip_existing: 是否跳过已存在的文档
            show_progress: 是否显示进度条

        Returns:
            实际插入的文档数量
        """
        if not documents:
            return 0

        total_docs = len(documents)
        logger.info(f"Starting batch upsert: {total_docs} documents, batch_size={batch_size}")

        skipped = 0
        if skip_existing:
            existing_ids = self._get_existing_ids([doc.id for doc in documents])
            skipped = len([doc for doc in documents if doc.id in existing_ids])
            documents = [doc for doc in documents if doc.id not in existing_ids]
            if skipped > 0:
                logger.info(f"Skipped {skipped} existing documents")
            if not documents:
                return 0

        inserted = 0
        total_batches = (len(documents) + batch_size - 1) // batch_size
        iterator = range(0, len(documents), batch_size)

        if show_progress:
            try:
                from tqdm import tqdm
                iterator = tqdm(iterator, desc="Upserting", total=total_batches, unit="batch")
            except ImportError:
                logger.debug("tqdm not installed, progress bar disabled")

        for i in iterator:
            batch = documents[i:i + batch_size]
            self.upsert(batch)
            inserted += len(batch)

        logger.info(f"Batch upsert completed: {inserted} inserted, {skipped} skipped")
        return inserted

    def _get_existing_ids(self, candidate_ids: List[str]) -> set:
        """获取已存在的文档 ID 集合"""
        if self._table is None:
            return set()

        all_ids = set(self.get_all_ids())
        return all_ids.intersection(candidate_ids)

    def migrate_to(
        self,
        target,
        batch_size: int = 100,
        skip_existing: bool = True,
        show_progress: bool = True,
    ) -> int:
        """
        将当前表的所有数据迁移到目标 retriever

        Args:
            target: 目标 retriever
            batch_size: 批处理大小
            skip_existing: 是否跳过已存在的文档
            show_progress: 是否显示进度条

        Returns:
            迁移的文档数量
        """
        all_ids = self.get_all_ids()
        if not all_ids:
            logger.info("No documents to migrate")
            return 0

        total = len(all_ids)
        logger.info(f"Starting migration: {total} documents")

        migrated = 0
        iterator = range(0, total, batch_size)

        if show_progress:
            try:
                from tqdm import tqdm
                iterator = tqdm(
                    iterator,
                    desc="Migrating",
                    total=(total + batch_size - 1) // batch_size,
                    unit="batch",
                )
            except ImportError:
                pass

        for i in iterator:
            batch_ids = all_ids[i:i + batch_size]
            documents = self.get(ids=batch_ids)
            if documents:
                migrated += target.upsert_batch(
                    documents,
                    batch_size=batch_size,
                    skip_existing=skip_existing,
                    show_progress=False,
                )

        logger.info(f"Migration completed: {migrated} documents migrated")
        return migrated

    def __repr__(self) -> str:
        extra = f", extra_fields={list(self._extra_fields)}" if self._extra_fields else ""
        return (
            f"LanceRetriever("
            f"table={self.table_name!r}, "
            f"count={self.count()}"
            f"{extra})"
        )
