import asyncio
import os
import sys
import types
from functools import lru_cache
from pathlib import Path
from typing import List, Dict, Any

import numpy as np
from minirag import MiniRAG, QueryParam, minirag as minirag_mod
from minirag.base import BaseKVStorage, BaseVectorStorage, BaseGraphStorage
from minirag.utils import wrap_embedding_func_with_attrs, compute_mdhash_id
from sentence_transformers import SentenceTransformer

from .config import PROJECT_ROOT, QA_PATH, QA_REPORT_PATH

# 环境设置：关闭实体抽取，避免额外依赖
os.environ.setdefault("MINIRAG_DISABLE_ENTITY_EXTRACT", "1")

# --------- 为 pip 版 minirag 注入轻量存储实现，避免缺失 kg 模块 ---------
class _JsonKVStorage(BaseKVStorage):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.storage: Dict[str, dict] = {}

    async def all_keys(self):
        return list(self.storage.keys())

    async def get_by_id(self, id: str):
        return self.storage.get(id)

    async def get_by_ids(self, ids, fields=None):
        out = []
        for i in ids:
            v = self.storage.get(i)
            if v and fields:
                v = {k: v[k] for k in fields if k in v}
            out.append(v)
        return out

    async def filter_keys(self, data):
        return {k for k in data if k not in self.storage}

    async def upsert(self, data):
        self.storage.update(data)

    async def drop(self):
        self.storage.clear()

    async def index_done_callback(self):
        return

    async def query_done_callback(self):
        return


class _NanoVectorDBStorage(BaseVectorStorage):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.storage: Dict[str, Dict[str, Any]] = {}
        self.embeddings: Dict[str, np.ndarray] = {}

    async def query(self, query: str, top_k: int):
        if not self.embeddings:
            return []
        q_embs = await self.embedding_func([query])
        q_emb = np.array(q_embs[0], dtype=np.float32)
        sims = []
        for k, emb in self.embeddings.items():
            denom = (np.linalg.norm(emb) * np.linalg.norm(q_emb)) or 1e-6
            sims.append((k, float(np.dot(emb, q_emb) / denom)))
        sims.sort(key=lambda x: x[1], reverse=True)
        res = []
        for k, score in sims[:top_k]:
            item = dict(self.storage.get(k, {}))
            item.update({"id": k, "score": score})
            res.append(item)
        return res

    async def upsert(self, data):
        texts = [v.get("content", "") for v in data.values()]
        embs = await self.embedding_func(texts)
        if isinstance(embs, np.ndarray):
            embs = list(embs)
        for (k, v), emb in zip(data.items(), embs):
            self.storage[k] = v
            self.embeddings[k] = np.array(emb, dtype=np.float32)

    async def index_done_callback(self):
        return

    async def query_done_callback(self):
        return


class _NetworkXStorage(BaseGraphStorage):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.nodes: Dict[str, dict] = {}
        self.edges: Dict[tuple, dict] = {}

    async def get_types(self):
        return [], []

    async def has_node(self, node_id: str):
        return node_id in self.nodes

    async def has_edge(self, source_node_id: str, target_node_id: str):
        return (source_node_id, target_node_id) in self.edges

    async def node_degree(self, node_id: str):
        return 0

    async def edge_degree(self, src_id: str, tgt_id: str):
        return 0

    async def get_node(self, node_id: str):
        return self.nodes.get(node_id)

    async def get_edge(self, source_node_id: str, target_node_id: str):
        return self.edges.get((source_node_id, target_node_id))

    async def get_node_edges(self, source_node_id: str):
        return []

    async def upsert_node(self, node_id: str, node_data: dict[str, str]):
        self.nodes[node_id] = node_data

    async def upsert_edge(self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]):
        self.edges[(source_node_id, target_node_id)] = edge_data

    async def delete_node(self, node_id: str):
        self.nodes.pop(node_id, None)
        for k in list(self.edges.keys()):
            if k[0] == node_id or k[1] == node_id:
                self.edges.pop(k, None)

    async def embed_nodes(self, algorithm: str):
        return np.zeros((0,)), []

    async def index_done_callback(self):
        return

    async def query_done_callback(self):
        return


class _JsonDocStatusStorage(_JsonKVStorage):
    pass


_simple_module = types.ModuleType("minirag.simple_storage")
_simple_module.JsonKVStorage = _JsonKVStorage
_simple_module.NanoVectorDBStorage = _NanoVectorDBStorage
_simple_module.NetworkXStorage = _NetworkXStorage
_simple_module.JsonDocStatusStorage = _JsonDocStatusStorage
sys.modules["minirag.simple_storage"] = _simple_module

minirag_mod.STORAGES.update({
    "NetworkXStorage": "minirag.simple_storage",
    "JsonKVStorage": "minirag.simple_storage",
    "NanoVectorDBStorage": "minirag.simple_storage",
    "JsonDocStatusStorage": "minirag.simple_storage",
})
# -------------------------------------------------------------------------

# 模型与工作目录路径
MODEL_DIR = (PROJECT_ROOT.parent / "minirag" / "minirag" / "models" / "bge-small-zh-v1.5").resolve()
WORKDIR = (PROJECT_ROOT / "minirag_cache").resolve()
WORKDIR.mkdir(parents=True, exist_ok=True)

# 预加载 QA
def _load_qas(path: Path) -> List[Dict[str, Any]]:
    return __import__("json").loads(path.read_text(encoding="utf-8"))


def _build_embedder():
    model = SentenceTransformer(str(MODEL_DIR), device="cpu")
    emb_dim = model.get_sentence_embedding_dimension()

    @wrap_embedding_func_with_attrs(embedding_dim=emb_dim, max_token_size=512)
    async def embed(texts):
        if isinstance(texts, str):
            texts = [texts]
        embs = model.encode(texts, normalize_embeddings=True, convert_to_numpy=True)
        return embs.astype(np.float32)

    return embed


@lru_cache(maxsize=1)
@lru_cache(maxsize=1)
def _embedder_cached():
    return _build_embedder()


def _build_rag_for(path: Path, workdir: Path):
    qas = _load_qas(path)
    embed = _embedder_cached()
    rag = MiniRAG(
        working_dir=str(workdir),
        embedding_func=embed,
        chunk_token_size=1200,      # 问+答一条，不再细切
        chunk_overlap_token_size=0,
        llm_model_func=lambda *a, **k: "",
        log_level="WARNING",
    )
    chunks = []
    id_to_qa = {}
    for qa in qas:
        chunk_text = f"Q{qa.get('id')}：{qa.get('question','')}\nA：{qa.get('answer','')}"
        cid = compute_mdhash_id(chunk_text, prefix="chunk-")
        chunks.append(chunk_text)
        id_to_qa[cid] = qa

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(rag.ainsert(chunks))
    loop.close()
    return rag, id_to_qa


@lru_cache(maxsize=1)
def _rag_bundle_common():
    workdir = WORKDIR / "common"
    workdir.mkdir(parents=True, exist_ok=True)
    return _build_rag_for(QA_PATH, workdir)


@lru_cache(maxsize=1)
def _rag_bundle_report():
    workdir = WORKDIR / "report"
    workdir.mkdir(parents=True, exist_ok=True)
    return _build_rag_for(QA_REPORT_PATH, workdir)


REPORT_KEYWORDS = [
    # 报告/证书类
    "检测", "检测报告", "检验", "检验报告", "质检", "质检报告", "第三方",
    "报告", "报告编号", "编号", "证书", "证明", "盖章", "章", "资质", "认证",
    # SDS / MSDS / COA
    "MSDS", "SDS", "安全数据单", "安全技术说明书", "COA", "COC", "CMA", "CNAS",
    # 指标/结果表达
    "成分", "成分表", "配方", "含量", "浓度", "数值", "指标", "限值", "限量",
    "合格", "达标", "超标", "未检出", "检出限", "ppm", "mg/kg",
    # 重点关注物质/风险词
    "化学", "毒", "有毒", "安全性", "刺激", "过敏", "致敏",
    "汞", "砷", "铅", "镉", "甲醇", "甲醛", "重金属",
]


def _hit_report(query: str):
    q = query.lower()
    return [kw for kw in REPORT_KEYWORDS if kw.lower() in q]


def _query_rag(rag, id_to_qa, query: str, limit: int, source: str):
    async def _search():
        results = await rag.chunks_vdb.query(query, top_k=limit)
        out = []
        for r in results:
            qa = id_to_qa.get(r.get("id"))
            if not qa:
                continue
            out.append({
                "id": qa.get("id"),
                "question": qa.get("question", ""),
                "answer": qa.get("answer", ""),
                "score": r.get("score", 0.0),
                "source": source,
            })
        return out

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(_search())
    finally:
        loop.close()


def search_rag(query: str, limit: int = 5) -> List[Dict[str, str]]:
    """
    使用 minirag 检索，返回 question/answer 列表（无 id，供简单问答）。
    命中检测关键词时，同时查检测库。
    """
    hits = _hit_report(query)
    rag_c, map_c = _rag_bundle_common()
    common = _query_rag(rag_c, map_c, query, limit, "common")
    report = []
    if hits:
        rag_r, map_r = _rag_bundle_report()
        report = _query_rag(rag_r, map_r, query, limit, "report")

    merged = []
    seen = set()
    for item in report + common:
        if item["id"] in seen:
            continue
        seen.add(item["id"])
        merged.append({"question": item["question"], "answer": item["answer"], "source": item["source"]})
        if len(merged) >= limit:
            break
    print(f"[RAG][faq_chat] used={'both' if hits else 'common'} kw={hits} q='{query}'")
    return merged


def search_rag_full(query: str, limit: int = 10) -> List[Dict[str, Any]]:
    """
    返回带 id / question / answer 的列表，供 FAQ 阶段展示。
    命中检测关键词时：检测库+常用库合并，检测结果优先。
    """
    hits = _hit_report(query)
    rag_c, map_c = _rag_bundle_common()
    common = _query_rag(rag_c, map_c, query, limit, "common")
    report = []
    if hits:
        rag_r, map_r = _rag_bundle_report()
        report = _query_rag(rag_r, map_r, query, limit, "report")

    merged = []
    seen = set()
    for item in report + common:
        if item["id"] in seen:
            continue
        seen.add(item["id"])
        merged.append({
            "id": item["id"],
            "question": item["question"],
            "answer": item["answer"],
            "source": item["source"],
            "score": item.get("score", 0.0),
        })
        if len(merged) >= limit:
            break

    debug = {"used_index": "both" if hits else "common", "hit_keywords": hits}
    print(f"[RAG][faq_search] used={debug['used_index']} kw={hits} q='{query}'")
    return merged, debug