import json
import os
import threading
from pathlib import Path
from typing import Any

import faiss
import numpy as np
import torch
import torch.nn.functional as F
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer


app = FastAPI(title="Bosun Knowledge Service", version="1.0.0")


class ClassifyRequest(BaseModel):
    query: str


class SearchRequest(BaseModel):
    query: str
    space_key: str = "general"
    top_k: int = 5


class IndexItem(BaseModel):
    id: int
    title: str
    content: str
    chunk_index: int = 0
    source_type: str = ""
    source_id: str = ""
    knowledge_space_key: str = "general"
    knowledge_space_name: str = "通用知识库"


class RebuildIndexRequest(BaseModel):
    index_dir: str
    items: list[IndexItem] = Field(default_factory=list)


class KnowledgeEngine:
    def __init__(self) -> None:
        self.device = "cuda" if os.getenv("KNOWLEDGE_DEVICE", "cpu") == "cuda" and torch.cuda.is_available() else "cpu"
        self.classifier_encoder_name = os.getenv("KNOWLEDGE_BERT_MODEL", "bert-base-chinese")
        self.classifier_model_name = os.getenv("KNOWLEDGE_BERT_CLASSIFIER_MODEL", "")
        self.embedding_model_name = os.getenv("KNOWLEDGE_BGE_MODEL", "BAAI/bge-small-zh-v1.5")
        self.lock = threading.RLock()
        self.classifier_tokenizer = None
        self.classifier_encoder = None
        self.classifier_head_tokenizer = None
        self.classifier_head_model = None
        self.embedding_tokenizer = None
        self.embedding_model = None
        self.index = None
        self.metadata: list[dict[str, Any]] = []
        self.index_dir = Path(os.getenv("KNOWLEDGE_INDEX_DIR", Path(__file__).resolve().parent / "data" / "faiss"))
        self.index_path = self.index_dir / "knowledge.index"
        self.meta_path = self.index_dir / "knowledge_meta.json"
        self.label_descriptions = {
            "faq": "标准业务问答、高频固定问题、制度规则查询、功能操作说明",
            "document": "需要查阅资料、归纳文档、分析内容、总结合同与报告的复杂问题",
            "smalltalk": "闲聊、故事、笑话、天气、娱乐、无业务价值的聊天问题",
            "out_of_scope": "股票、医疗、法律、旅游、影视等超出企业知识库范围的问题",
            "invalid": "信息过少、语义不完整、无法判断意图的无效提问",
        }
        self._load_index_if_exists()

    def health(self) -> dict[str, Any]:
        return {
            "device": self.device,
            "classifier_encoder": self.classifier_encoder_name,
            "classifier_model": self.classifier_model_name or "prototype-bert",
            "embedding_model": self.embedding_model_name,
            "index_dir": str(self.index_dir),
            "indexed_items": len(self.metadata),
        }

    def classify(self, query: str) -> dict[str, Any]:
        query = (query or "").strip()
        if len(query) < 2:
            return {"intent": "invalid", "score": 1.0, "reason": "query_too_short"}

        if self.classifier_model_name:
            try:
                return self._classify_with_sequence_model(query)
            except Exception:
                pass

        try:
            return self._classify_with_prototypes(query)
        except Exception:
            pass

        rule_result = self._rule_classify(query)
        if rule_result:
            return rule_result

        return {"intent": "document", "score": 0.0, "reason": "bert_classifier_unavailable"}

    def _rule_classify(self, query: str) -> dict[str, Any] | None:
        text = query.strip().lower()
        if not text:
            return {"intent": "invalid", "score": 1.0, "reason": "empty_query"}

        smalltalk_keywords = [
            "讲个故事",
            "小故事",
            "笑话",
            "闲聊",
            "聊天",
            "夸夸我",
            "今天天气",
            "天气",
            "星座",
            "你好",
            "hi",
            "hello",
            "你是谁",
            "你能做什么",
            "你是怎么工作的",
        ]
        if any(keyword in text for keyword in smalltalk_keywords):
            return {"intent": "smalltalk", "score": 0.98, "reason": "rule_smalltalk"}

        out_of_scope_keywords = [
            "股票",
            "彩票",
            "电影",
            "明星",
            "旅游",
            "菜谱",
            "医学",
            "法律咨询",
            "写诗",
            "翻译成英文",
        ]
        if any(keyword in text for keyword in out_of_scope_keywords):
            return {"intent": "out_of_scope", "score": 0.98, "reason": "rule_out_of_scope"}

        document_keywords = [
            "总结",
            "梳理",
            "分析",
            "对比",
            "归纳",
            "提纲",
            "解读",
            "起草",
            "生成",
            "合同",
            "报告",
            "方案",
            "条款",
            "根据资料",
            "根据文档",
            "根据知识库",
            "整理",
            "审批要点",
        ]
        if any(keyword in text for keyword in document_keywords):
            return {"intent": "document", "score": 0.97, "reason": "rule_document"}

        faq_keywords = [
            "如何",
            "怎么",
            "哪里",
            "在哪",
            "是否",
            "有没有",
            "可以",
            "支持",
            "密码",
            "登录",
            "佣金",
            "规则",
            "流程",
            "审批",
            "上传",
            "删除",
            "新建",
        ]
        if any(keyword in text for keyword in faq_keywords):
            return {"intent": "faq", "score": 0.96, "reason": "rule_faq"}

        if len(query) <= 18:
            return {"intent": "faq", "score": 0.75, "reason": "rule_short_query"}

        return None

    def search(self, query: str, space_key: str, top_k: int) -> list[dict[str, Any]]:
        query = (query or "").strip()
        if not query:
            return []
        with self.lock:
            if self.index is None or not self.metadata:
                self._load_index_if_exists()
            if self.index is None or not self.metadata:
                return []

        query_vec = self._encode_embedding([query])[0].astype("float32")
        query_vec = np.expand_dims(query_vec, axis=0)
        fetch_k = min(len(self.metadata), max(top_k * 8, 32))
        scores, indices = self.index.search(query_vec, fetch_k)
        items: list[dict[str, Any]] = []
        for score, idx in zip(scores[0].tolist(), indices[0].tolist()):
            if idx < 0 or idx >= len(self.metadata):
                continue
            meta = self.metadata[idx]
            if not self._match_space(space_key, meta.get("knowledge_space_key", "general")):
                continue
            items.append(
                {
                    **meta,
                    "snippet": self._build_snippet(meta.get("content", ""), query),
                    "score": float(score),
                }
            )
            if len(items) >= top_k:
                break
        return items

    def rebuild_index(self, index_dir: str, items: list[IndexItem]) -> dict[str, Any]:
        target_dir = Path(index_dir or self.index_dir)
        target_dir.mkdir(parents=True, exist_ok=True)
        records = [item.model_dump() for item in items if item.content.strip()]

        if not records:
            empty_index = faiss.IndexFlatIP(384)
            faiss.write_index(empty_index, str(target_dir / "knowledge.index"))
            (target_dir / "knowledge_meta.json").write_text("[]", encoding="utf-8")
            with self.lock:
                self.index_dir = target_dir
                self.index_path = target_dir / "knowledge.index"
                self.meta_path = target_dir / "knowledge_meta.json"
                self.index = empty_index
                self.metadata = []
            return {"indexed_items": 0}

        embeddings = self._encode_embedding([item["content"] for item in records]).astype("float32")
        dim = int(embeddings.shape[1])
        index = faiss.IndexFlatIP(dim)
        index.add(embeddings)
        faiss.write_index(index, str(target_dir / "knowledge.index"))
        (target_dir / "knowledge_meta.json").write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")
        with self.lock:
            self.index_dir = target_dir
            self.index_path = target_dir / "knowledge.index"
            self.meta_path = target_dir / "knowledge_meta.json"
            self.index = index
            self.metadata = records
        return {"indexed_items": len(records)}

    def _load_index_if_exists(self) -> None:
        with self.lock:
            if self.index_path.exists() and self.meta_path.exists():
                self.index = faiss.read_index(str(self.index_path))
                self.metadata = json.loads(self.meta_path.read_text(encoding="utf-8"))

    def _classify_with_sequence_model(self, query: str) -> dict[str, Any]:
        if self.classifier_head_tokenizer is None or self.classifier_head_model is None:
            self.classifier_head_tokenizer = AutoTokenizer.from_pretrained(self.classifier_model_name)
            self.classifier_head_model = AutoModelForSequenceClassification.from_pretrained(self.classifier_model_name).to(self.device)
            self.classifier_head_model.eval()
        encoded = self.classifier_head_tokenizer(query, return_tensors="pt", truncation=True, max_length=256).to(self.device)
        with torch.no_grad():
            logits = self.classifier_head_model(**encoded).logits
            probs = F.softmax(logits, dim=-1)[0].detach().cpu().numpy()
        labels = self.classifier_head_model.config.id2label or {}
        best_idx = int(np.argmax(probs))
        intent = str(labels.get(best_idx, "document")).lower()
        if intent not in self.label_descriptions:
            intent = "document"
        return {"intent": intent, "score": float(probs[best_idx]), "reason": "sequence_classifier"}

    def _classify_with_prototypes(self, query: str) -> dict[str, Any]:
        texts = [query] + list(self.label_descriptions.values())
        vectors = self._encode_classifier(texts)
        query_vec = vectors[0]
        label_vecs = vectors[1:]
        scores = np.dot(label_vecs, query_vec)
        labels = list(self.label_descriptions.keys())
        best_idx = int(np.argmax(scores))
        return {
            "intent": labels[best_idx],
            "score": float(scores[best_idx]),
            "reason": "prototype_bert",
        }

    def _encode_classifier(self, texts: list[str]) -> np.ndarray:
        if self.classifier_tokenizer is None or self.classifier_encoder is None:
            self.classifier_tokenizer = AutoTokenizer.from_pretrained(self.classifier_encoder_name)
            self.classifier_encoder = AutoModel.from_pretrained(self.classifier_encoder_name).to(self.device)
            self.classifier_encoder.eval()
        return self._encode(texts, self.classifier_tokenizer, self.classifier_encoder)

    def _encode_embedding(self, texts: list[str]) -> np.ndarray:
        if self.embedding_tokenizer is None or self.embedding_model is None:
            self.embedding_tokenizer = AutoTokenizer.from_pretrained(self.embedding_model_name)
            self.embedding_model = AutoModel.from_pretrained(self.embedding_model_name).to(self.device)
            self.embedding_model.eval()
        return self._encode(texts, self.embedding_tokenizer, self.embedding_model)

    def _encode(self, texts: list[str], tokenizer: Any, model: Any) -> np.ndarray:
        batches: list[np.ndarray] = []
        batch_size = 16
        for start in range(0, len(texts), batch_size):
            batch = texts[start : start + batch_size]
            encoded = tokenizer(
                batch,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt",
            ).to(self.device)
            with torch.no_grad():
                output = model(**encoded)
            pooled = self._mean_pool(output.last_hidden_state, encoded["attention_mask"])
            pooled = F.normalize(pooled, p=2, dim=1)
            batches.append(pooled.detach().cpu().numpy())
        return np.vstack(batches)

    @staticmethod
    def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
        summed = torch.sum(last_hidden_state * mask, dim=1)
        counts = torch.clamp(mask.sum(dim=1), min=1e-9)
        return summed / counts

    @staticmethod
    def _match_space(selected: str, current: str) -> bool:
        selected = (selected or "general").strip().lower()
        current = (current or "general").strip().lower()
        if selected in {"", "general", "all"}:
            return True
        return current in {selected, "general"}

    @staticmethod
    def _build_snippet(content: str, query: str) -> str:
        text = " ".join(content.split())
        if not text:
            return ""
        terms = [term for term in query.replace("，", " ").replace("。", " ").split() if term]
        hit = next((term for term in terms if term in text), "")
        if not hit:
            return text[:180]
        pos = text.find(hit)
        start = max(0, pos - 60)
        end = min(len(text), pos + 120)
        return text[start:end]


engine = KnowledgeEngine()


@app.get("/health")
def health() -> dict[str, Any]:
    return {"status": "ok", "data": engine.health()}


@app.post("/classify")
def classify(req: ClassifyRequest) -> dict[str, Any]:
    try:
        return {"data": engine.classify(req.query)}
    except Exception as exc:  # pragma: no cover - runtime guard
        raise HTTPException(status_code=500, detail=f"classify failed: {exc}") from exc


@app.post("/search")
def search(req: SearchRequest) -> dict[str, Any]:
    try:
        return {"data": {"items": engine.search(req.query, req.space_key, max(1, min(req.top_k, 20))) }}
    except Exception as exc:  # pragma: no cover - runtime guard
        raise HTTPException(status_code=500, detail=f"search failed: {exc}") from exc


@app.post("/index/rebuild")
def rebuild_index(req: RebuildIndexRequest) -> dict[str, Any]:
    try:
        return {"data": engine.rebuild_index(req.index_dir, req.items)}
    except Exception as exc:  # pragma: no cover - runtime guard
        raise HTTPException(status_code=500, detail=f"rebuild failed: {exc}") from exc