"""
词频统计（jieba 分词版）：流程与 词频.py 相同，第 2 步改用 jieba 分词。

  1. 随机样本 → LLM 归纳专有名词与停用词（与 spaCy 版共用 voc_terms.json）
  2. jieba 全量 content 分词 + 词频 → output/word_freq.csv

用法::

  python3 词频_jieba.py
  python3 词频_jieba.py --skip-llm
"""
from __future__ import annotations

import argparse
import json
import logging
import re
import sys
from collections import Counter
from pathlib import Path
from typing import List, Sequence, Set, Tuple

from 词频 import (
    OUTPUT_DIR,
    SAMPLE_SEED,
    SAMPLE_SIZE,
    STRUCTURED_DB,
    TERMS_JSON,
    WORD_FREQ_CSV,
    _apply_product_terms,
    _finalize_term_lists,
    _is_pure_number,
    _latest_job,
    _load_contents,
    _load_spacy,
    _load_terms_json,
    _merge_word_forms,
    _overlaps_span,
    _resolve_api_key,
    _resolve_source_path,
    _sample_reviews,
    _save_word_freq,
    _step1_extract_terms,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger("voc_wordfreq_jieba")

# 英文词或中文词均保留
_TOKEN_OK = re.compile(r"[a-z\u4e00-\u9fff]", re.I)


def _prepare_jieba(product_terms: Sequence[str]) -> None:
    import jieba

    for term in product_terms:
        t = (term or "").strip()
        if t:
            jieba.add_word(t)


def _tokenize_jieba(
    text: str,
    protected_spans: Sequence[Tuple[int, int]],
    stop_set: Set[str],
) -> List[str]:
    import jieba

    tokens: List[str] = []
    for word, start, end in jieba.tokenize(text):
        if _overlaps_span(start, end, protected_spans):
            continue
        w = word.strip().lower()
        if not w or _is_pure_number(w) or w in stop_set:
            continue
        if not _TOKEN_OK.search(w):
            continue
        tokens.append(w)
    return tokens


def _build_word_freq(
    rows: Sequence[Tuple[int, str]],
    product_terms: List[str],
    stop_set: Set[str],
) -> Counter[str]:
    _prepare_jieba(product_terms)
    counter: Counter[str] = Counter()
    for _, text in rows:
        if not text.strip():
            continue
        lower = text.lower()
        protected = _apply_product_terms(lower, product_terms, counter, stop_set)
        for tok in _tokenize_jieba(text, protected, stop_set):
            counter[tok] += 1
    counter = Counter({k: v for k, v in counter.items() if not _is_pure_number(k)})
    nlp = _load_spacy()
    return _merge_word_forms(counter, nlp)


def run(*, skip_llm: bool = False) -> dict:
    if not STRUCTURED_DB.is_file():
        raise FileNotFoundError(f"缺少 {STRUCTURED_DB}")

    import sqlite3

    conn = sqlite3.connect(STRUCTURED_DB)
    try:
        job_id, industry, product_name, source_file = _latest_job(conn)
    finally:
        conn.close()

    csv_path = _resolve_source_path(source_file)
    rows = _load_contents(csv_path)
    logger.info(
        "job_id=%s product=%r 评论 %s 条，来源 %s（jieba 分词）",
        job_id,
        product_name,
        len(rows),
        csv_path.name,
    )

    if skip_llm:
        meta = _load_terms_json()
        product_terms, stop_set = _finalize_term_lists(
            list(meta.get("product_terms") or []),
            list(meta.get("stopwords_custom") or []),
            product_name,
        )
        logger.info("第 1 步跳过，复用 %s", TERMS_JSON)
    else:
        api_key = _resolve_api_key()
        if not api_key:
            raise RuntimeError("缺少 DASHSCOPE_API_KEY 或 .dashscope_key")
        samples = _sample_reviews(rows, n=SAMPLE_SIZE, seed=SAMPLE_SEED)
        meta = _step1_extract_terms(
            job_id=job_id,
            industry=industry,
            product_name=product_name,
            samples=samples,
            api_key=api_key,
        )
        product_terms, stop_set = _finalize_term_lists(
            meta["product_terms"],
            meta.get("stopwords_custom") or [],
            product_name,
        )

    logger.info("第 2 步：jieba 全量分词与词频统计")
    counter = _build_word_freq(rows, product_terms, stop_set)
    _save_word_freq(counter, WORD_FREQ_CSV)
    logger.info("已写入 %s", WORD_FREQ_CSV)

    return {
        "job_id": job_id,
        "product_name": product_name,
        "tokenizer": "jieba",
        "terms_json": str(TERMS_JSON),
        "word_freq_csv": str(WORD_FREQ_CSV),
        "unique_words": len(counter),
        "total_tokens": sum(counter.values()),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="VOC 词频统计（jieba 分词）")
    parser.add_argument(
        "--skip-llm",
        action="store_true",
        help="跳过第 1 步，使用 output/voc_terms.json",
    )
    args = parser.parse_args()
    result = run(skip_llm=args.skip_llm)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()