"""关键词检索 + 真实模型草稿。运行：python retrieval_qa.py；仅检索：加 --retrieve-only。"""
import argparse
import getpass
import json
import os
from pathlib import Path
import sys
import warnings

from model_reply_assistant import answer as generate


def load_documents(path):
    path = Path(path)
    if path.stat().st_size > 100_000:
        raise ValueError('资料文件超过本课 100 KB 限制。')
    docs = json.loads(path.read_text(encoding='utf-8-sig'))
    if not isinstance(docs, list) or not 1 <= len(docs) <= 100:
        raise ValueError('资料必须是 1–100 个片段的 JSON 数组。')
    seen = set()
    for doc in docs:
        if not isinstance(doc, dict) or set(doc) != {'id', 'title', 'text', 'keywords'}:
            raise ValueError('每个片段必须包含 id、title、text、keywords。')
        for field, limit in [('id', 80), ('title', 120), ('text', 2000)]:
            if not isinstance(doc[field], str) or not 1 <= len(doc[field].strip()) <= limit:
                raise ValueError('资料字段为空、类型错误或过长。')
        if doc['id'] in seen:
            raise ValueError('片段编号重复。')
        seen.add(doc['id'])
        terms = doc['keywords']
        if not isinstance(terms, list) or not 1 <= len(terms) <= 20 or any(not isinstance(t, str) or not 1 <= len(t.strip()) <= 40 for t in terms):
            raise ValueError('keywords 必须是 1–20 个非空短词。')
    return docs


def retrieve(question, documents):
    question = question.strip().casefold()
    if not 1 <= len(question) <= 200:
        raise ValueError('请输入 1–200 字的问题。')
    scored = []
    for doc in documents:
        matched = sorted({term.strip().casefold() for term in doc['keywords'] if term.strip().casefold() in question})
        if matched:
            scored.append((doc, matched))
    # shortcut: 小资料集按命中关键词数排序，出现同义漏检等实际问题时再评估语义检索。
    scored.sort(key=lambda item: (-len(item[1]), item[0]['id']))
    return [{'id': d['id'], 'title': d['title'], 'text': d['text'], 'matched': terms} for d, terms in scored[:3]]


def answer(question, api_key, documents):
    hits = retrieve(question, documents)
    if not hits:
        return {'status': 'missing', 'draft': '未检索到相关片段，请换个问法、补充资料或转交人工。', 'sources': []}
    selected = [{key: hit[key] for key in ('id', 'title', 'text')} for hit in hits]
    return generate(question, api_key, selected)


def main():
    parser = argparse.ArgumentParser(description='虚构资料检索问答：先查片段，再请求真实模型。')
    parser.add_argument('--retrieve-only', action='store_true', help='只看检索结果，不调用模型，不需要密钥')
    args = parser.parse_args()
    try:
        docs = load_documents(Path(__file__).with_name('retrieval_documents.json'))
        print('仅使用公开或虚构资料；模型调用可能计费，不自动发送回复。')
        question = input('你的问题（1–200 字）：')
        hits = retrieve(question, docs)
        print('检索片段（最多 3 条，分数仅为关键词命中数）：')
        print(json.dumps(hits, ensure_ascii=False, indent=2))
        if args.retrieve_only:
            print('仅检索完成：未调用模型。')
            return 0
        if not hits:
            print('资料不足：未调用模型。请换个问法、补充资料或转交人工。')
            return 0
        key = os.environ.get('DEEPSEEK_API_KEY')
        if not key:
            with warnings.catch_warnings():
                warnings.simplefilter('error', getpass.GetPassWarning)
                key = getpass.getpass('DeepSeek API Key（输入不回显）：')
        result = answer(question, key, docs)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        print('引用仅允许来自本次命中片段；逐条核对事实后再使用，程序不会发送回复。')
        return 0
    except RuntimeError as error:
        print('未完成：' + str(error), file=sys.stderr)
        return 1
    except (OSError, ValueError):
        print('未完成：请检查资料文件是否存在、JSON 格式与字段，以及输入和密钥格式。未交付草稿。', file=sys.stderr)
        return 1
    except getpass.GetPassWarning:
        print('未完成：当前终端不能隐藏密钥，请换交互终端或安全配置环境变量。', file=sys.stderr)
        return 1
    except (EOFError, KeyboardInterrupt):
        print('\n已取消。', file=sys.stderr)
        return 1


if __name__ == '__main__':
    sys.exit(main())