#!/usr/bin/env python3
"""RAG 问答模块：向量检索 + LLM 生成。

流程：
1. 将用户问题转为向量
2. 在 embeddings 表中做余弦相似度检索，取 TOP-K 相关文章
3. 将检索到的文章内容作为上下文，调用千问生成回答
4. 返回回答 + 引用来源（可追溯）
"""

from __future__ import annotations

import json
import os
import sqlite3
import urllib.error
import urllib.request
from pathlib import Path

import numpy as np

from dotenv import load_dotenv
from structurer import get_api_key, get_embeddings, DASHSCOPE_CHAT_URL

BASE_DIR = Path(__file__).resolve().parent
load_dotenv(BASE_DIR / ".env")


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """计算两个向量的余弦相似度。"""
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a, b) / (norm_a * norm_b))


def vector_search(conn: sqlite3.Connection, query_embedding: list[float], top_k: int = 5) -> list[dict]:
    """向量检索：在 embeddings 表中找到与查询向量最相似的文章。"""
    query_vec = np.array(query_embedding, dtype=np.float32)
    rows = conn.execute(
        """
        SELECT e.article_id, e.embedding, e.chunk_text,
               a.title, a.source_name, a.source_type, a.published_at, a.summary,
               x.result_json
        FROM embeddings e
        JOIN articles a ON a.id = e.article_id
        LEFT JOIN ai_analyses x ON x.id = (
            SELECT x2.id FROM ai_analyses x2
            WHERE x2.article_id = e.article_id AND x2.analysis_type = 'editorial'
            ORDER BY x2.updated_at DESC, x2.id DESC LIMIT 1
        )
        WHERE a.crawl_status = 'success'
        """
    ).fetchall()

    scored = []
    for row in rows:
        emb = json.loads(row["embedding"])
        emb_vec = np.array(emb, dtype=np.float32)
        score = cosine_similarity(query_vec, emb_vec)
        scored.append((score, row))
    scored.sort(key=lambda x: x[0], reverse=True)

    results = []
    seen_articles = set()
    for score, row in scored:
        if row["article_id"] in seen_articles:
            continue
        seen_articles.add(row["article_id"])
        analysis = json.loads(row["result_json"]) if row["result_json"] else {}
        results.append({
            "article_id": row["article_id"],
            "title": row["title"],
            "source": row["source_name"] or row["source_type"],
            "date": (row["published_at"] or "")[:10],
            "summary": analysis.get("summary") or row["summary"] or "",
            "key_points": analysis.get("key_points", []),
            "content": row["chunk_text"] or "",
            "score": score,
        })
        if len(results) >= top_k:
            break
    return results


def generate_answer(question: str, contexts: list[dict]) -> dict:
    """基于检索到的上下文，调用千问生成回答。"""
    api_key = get_api_key()
    model = os.getenv("QWEN_MODEL", "qwen-plus")

    # 构建上下文
    context_text = ""
    sources = []
    for i, ctx in enumerate(contexts, start=1):
        context_text += f"\n--- 来源{i} ---\n标题：{ctx['title']}\n摘要：{ctx['summary']}\n要点：{'; '.join(ctx['key_points'][:3])}\n原文片段：{ctx['content']}\n"
        sources.append({
            "title": ctx["title"],
            "source": ctx["source"],
            "date": ctx["date"],
        })

    prompt = f"""你是餐饮行业知识库的问答助手。根据以下知识库中的情报内容回答用户问题。

要求：
1. 回答必须基于提供的来源内容，不要编造信息
2. 如果来源内容不足以回答问题，明确说明"知识库中暂无直接相关信息"
3. 回答要结构化、简洁、有经营参考价值
4. 用中文回答

知识库来源：
{context_text}

用户问题：{question}"""

    payload = json.dumps({
        "model": model,
        "messages": [
            {"role": "system", "content": "你是餐饮行业知识库问答助手，基于库内情报给出有出处的专业回答。"},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.3,
    }, ensure_ascii=False).encode("utf-8")

    req = urllib.request.Request(
        DASHSCOPE_CHAT_URL,
        data=payload,
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=90) as response:
            body = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:500]
        raise RuntimeError(f"千问接口返回 {exc.code}：{detail}") from exc

    answer = body["choices"][0]["message"]["content"].strip()
    return {
        "answer": answer,
        "sources": sources,
        "contexts": contexts,
    }


def generate_answer_stream(question: str, contexts: list[dict]):
    """流式生成回答，逐 token yield。"""
    api_key = get_api_key()
    model = os.getenv("QWEN_MODEL", "qwen-plus")

    context_text = ""
    sources = []
    for i, ctx in enumerate(contexts, start=1):
        context_text += f"\n--- 来源{i} ---\n标题：{ctx['title']}\n摘要：{ctx['summary']}\n要点：{'; '.join(ctx['key_points'][:3])}\n原文片段：{ctx['content']}\n"
        sources.append({
            "title": ctx["title"],
            "source": ctx["source"],
            "date": ctx["date"],
        })

    prompt = f"""你是餐饮行业知识库的问答助手。根据以下知识库中的情报内容回答用户问题。

要求：
1. 回答必须基于提供的来源内容，不要编造信息
2. 如果来源内容不足以回答问题，明确说明"知识库中暂无直接相关信息"
3. 回答要结构化、简洁、有经营参考价值
4. 用中文回答，使用 Markdown 格式

知识库来源：
{context_text}

用户问题：{question}"""

    payload = json.dumps({
        "model": model,
        "messages": [
            {"role": "system", "content": "你是餐饮行业知识库问答助手，基于库内情报给出有出处的专业回答。使用 Markdown 格式输出。"},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.3,
        "stream": True,
        "stream_options": {"include_usage": False},
    }, ensure_ascii=False).encode("utf-8")

    req = urllib.request.Request(
        DASHSCOPE_CHAT_URL,
        data=payload,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
        },
        method="POST",
    )
    try:
        resp = urllib.request.urlopen(req, timeout=120)
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:500]
        raise RuntimeError(f"千问接口返回 {exc.code}：{detail}") from exc

    buffer = b""
    while True:
        chunk = resp.read(4096)
        if not chunk:
            break
        buffer += chunk
        while b"\n" in buffer:
            line, buffer = buffer.split(b"\n", 1)
            line = line.strip()
            if not line or line == b"data: [DONE]":
                continue
            if line.startswith(b"data:"):
                try:
                    data = json.loads(line[5:].strip())
                    delta = data.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content")
                    if content:
                        yield content
                except (json.JSONDecodeError, IndexError, KeyError):
                    pass
    resp.close()

    yield json.dumps({"__sources__": sources}, ensure_ascii=False)


def rag_qa(question: str, top_k: int = 5) -> dict:
    """完整的 RAG 问答流程：向量检索 → LLM 生成。"""
    db_path = Path(os.getenv("CIBANK_DB", str(BASE_DIR / "data" / "cibank.db")))
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row

    # 检查是否有向量数据
    count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
    if count == 0:
        # 无向量数据时使用关键词检索作为降级方案
        return keyword_fallback(conn, question)

    # 1. 将问题转为向量
    query_embeddings = get_embeddings([question])

    # 2. 向量检索
    contexts = vector_search(conn, query_embeddings[0], top_k)

    if not contexts:
        conn.close()
        return {
            "answer": "知识库中暂无与您问题相关的内容。请先运行采集和分析流程入库更多情报。",
            "sources": [],
            "contexts": [],
        }

    # 3. LLM 生成回答
    result = generate_answer(question, contexts)
    conn.close()
    return result


def rag_qa_stream(question: str, top_k: int = 5):
    """流式 RAG 问答：先检索，再流式生成。yield (type, data) 元组。"""
    db_path = Path(os.getenv("CIBANK_DB", str(BASE_DIR / "data" / "cibank.db")))
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row

    count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
    if count == 0:
        conn.close()
        yield ("error", "知识库中暂无向量数据，请先运行采集和分析流程。")
        return

    query_embeddings = get_embeddings([question])
    contexts = vector_search(conn, query_embeddings[0], top_k)
    conn.close()

    if not contexts:
        yield ("error", "知识库中暂无与您问题相关的内容。")
        return

    sources = [
        {"title": ctx["title"], "source": ctx["source"], "date": ctx["date"]}
        for ctx in contexts
    ]
    yield ("sources", sources)

    for token in generate_answer_stream(question, contexts):
        yield ("token", token)


def keyword_fallback(conn: sqlite3.Connection, question: str) -> dict:
    """无向量数据时的关键词检索降级方案。"""
    rows = conn.execute(
        """
        SELECT a.id, a.title, a.source_name, a.source_type, a.published_at, a.summary,
               a.content, x.result_json
        FROM articles a
        LEFT JOIN ai_analyses x ON x.article_id = a.id AND x.analysis_type = 'editorial'
        WHERE a.crawl_status = 'success'
          AND (a.title LIKE ? OR a.summary LIKE ? OR a.content LIKE ?)
        ORDER BY a.published_at DESC
        LIMIT 5
        """,
        (f"%{question}%", f"%{question}%", f"%{question}%"),
    ).fetchall()

    if not rows:
        conn.close()
        return {
            "answer": "知识库中暂无与您问题相关的内容。请先运行采集和分析流程入库更多情报。",
            "sources": [],
            "contexts": [],
        }

    contexts = []
    for row in rows:
        analysis = json.loads(row["result_json"]) if row["result_json"] else {}
        contexts.append({
            "article_id": row["id"],
            "title": row["title"],
            "source": row["source_name"] or row["source_type"],
            "date": (row["published_at"] or "")[:10],
            "summary": analysis.get("summary") or row["summary"] or "",
            "key_points": analysis.get("key_points", []),
            "content": row["content"][:3000] if row["content"] else "",
            "score": 0.5,
        })
    conn.close()
    return generate_answer(question, contexts)


if __name__ == "__main__":
    import sys
    question = " ".join(sys.argv[1:]) or "现在开茶饮店还有机会吗？"
    result = rag_qa(question)
    print(json.dumps(result, ensure_ascii=False, indent=2))