"""
词频统计：从最新结构化任务读取产品与 CSV，两步连跑。

  1. 随机 25 条 content → LLM 归纳「产品专有名词」与「Amazon/产品专属停用词」
  2. spaCy 全量 content 分词 + 词频 → output/word_freq.csv（专有名词按完整短语统计，不拆词）

用法::

  python3 词频词云.py
  python3 词频词云.py --skip-llm   # 复用 output/voc_terms.json，仅跑第 2 步
"""
from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import random
import re
import sqlite3
import sys
from collections import Counter
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Set, Tuple

from openai import OpenAI
from spacy.lang.en.stop_words import STOP_WORDS as EN_STOP_WORDS

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

PROJECT_ROOT = Path(__file__).resolve().parent
STRUCTURED_DB = PROJECT_ROOT / "voc_structured.sqlite"
OUTPUT_DIR = PROJECT_ROOT / "output"
TERMS_JSON = OUTPUT_DIR / "voc_terms.json"
WORD_FREQ_CSV = OUTPUT_DIR / "word_freq.csv"

DASHSCOPE_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
MODEL_NAME = os.environ.get("DASHSCOPE_MODEL", "qwen3.6-flash").strip()

SAMPLE_SIZE = 25
SAMPLE_SEED = 42


def _resolve_api_key() -> str:
    v = os.environ.get("DASHSCOPE_API_KEY", "").strip()
    if v:
        return v
    fp = os.environ.get("DASHSCOPE_API_KEY_FILE", "").strip()
    if fp:
        p = Path(fp).expanduser()
        if p.is_file():
            return p.read_text(encoding="utf-8").strip().strip('"').strip("'")
    local = PROJECT_ROOT / ".dashscope_key"
    if local.is_file():
        return local.read_text(encoding="utf-8").strip().strip('"').strip("'")
    return ""


def _strip_think(text: str) -> str:
    if not text:
        return text
    text = re.sub(
        r"<think>[\s\S]*?</think>", "", text, flags=re.IGNORECASE
    )
    text = re.sub(r"</?think>", "", text, flags=re.IGNORECASE)
    return text.strip()


def _latest_job(conn: sqlite3.Connection) -> Tuple[int, str, str, str]:
    row = conn.execute(
        """
        SELECT id, industry, product_name, source_file
        FROM analysis_jobs
        ORDER BY id DESC
        LIMIT 1
        """
    ).fetchone()
    if not row:
        raise RuntimeError(f"{STRUCTURED_DB} 中无 analysis_jobs 记录，请先跑结构化")
    return int(row[0]), str(row[1]), str(row[2]), str(row[3])


def _resolve_source_path(source_file: str) -> Path:
    p = Path(source_file)
    if p.is_file():
        return p.resolve()
    candidate = PROJECT_ROOT / source_file
    if candidate.is_file():
        return candidate.resolve()
    raise FileNotFoundError(f"找不到评论 CSV: {source_file}")


def _load_contents(csv_path: Path) -> List[Tuple[int, str]]:
    rows: List[Tuple[int, str]] = []
    with csv_path.open(encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames or "content" not in reader.fieldnames:
            raise ValueError(f"{csv_path} 缺少 content 列")
        for i, row in enumerate(reader, start=1):
            text = (row.get("content") or "").strip()
            if text:
                rows.append((i, text))
    if not rows:
        raise ValueError(f"{csv_path} 无有效 content")
    return rows


def _sample_reviews(
    rows: Sequence[Tuple[int, str]], *, n: int, seed: int
) -> List[Tuple[int, str]]:
    rng = random.Random(seed)
    k = min(n, len(rows))
    return rng.sample(list(rows), k)


def _build_terms_prompt(
    industry: str, product_name: str, samples: Sequence[Tuple[int, str]]
) -> Tuple[str, str]:
    blocks = []
    for row_id, text in samples:
        blocks.append(f"[R{row_id}]\n{text}")
    joined = "\n\n---\n\n".join(blocks)
    system = (
        "你是亚马逊美国站英文评论与电商文本分析专家。"
        "根据样本评论，列出两类词表，供后续英文分词与词频统计使用。"
    )
    user = f"""行业：{industry}
产品：{product_name}

以下为该产品 {len(samples)} 条随机评论（仅正文 content）：

{joined}

请输出两部分（可用 Markdown 标题、编号列表或逗号分隔，不必是 JSON）：

## 产品专有名词（product_terms）
列出该产品评论中可能出现的**全部**专有表达，尽量合理扩展，包括但不限于：
成分/品类、宠物品种与病种、剂型包装、用法场景、Amazon 物流售后相关但**属于本产品语境**的词等。
多词短语请保留完整形式（如 turkey tail mushroom powder）；后续词频**只统计完整短语**，不会拆成单词分别计数。
**不要**把下面「专属停用词」里的词放进本列表。

## 专属停用词（stopwords）
列出 Amazon 美国站电商通用词，以及与本产品强相关、词频高但**分析价值低**的词（如泛化评价词、过泛产品描述词等）。
多词短语可保留。
注意：不要把产品名「{product_name}」**整句**列为停用词；也不要把其中的**实词/品类成分**（如 turkey、tail、mushroom）列为停用词。
产品名拆词里**无分析价值的虚词/功能词**（如 for、the、a、of、and）可以且应列入停用词。

每条一行或逗号分隔即可。"""
    return system, user


def _call_llm(system: str, user: str, api_key: str) -> str:
    client = OpenAI(api_key=api_key, base_url=DASHSCOPE_BASE_URL)
    extra_body: Dict[str, Any] = {}
    if MODEL_NAME.lower().startswith(("qwen3.6", "qwen3.5", "qwen3")):
        extra_body["enable_thinking"] = False
    resp = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.3,
        **({"extra_body": extra_body} if extra_body else {}),
    )
    msg = resp.choices[0].message
    text = msg.content or getattr(msg, "reasoning_content", None) or ""
    if not text.strip():
        raise RuntimeError("LLM 返回为空")
    return _strip_think(text)


def _normalize_term(s: str) -> str:
    s = re.sub(r"\s+", " ", (s or "").strip().lower())
    return s.strip(".,;:\"'""''`-–—()[]{}")


def _split_list_items(line: str) -> List[str]:
    line = re.sub(r"^[\s\d\.\)\-•\*]+", "", line.strip())
    if not line or line.startswith("#"):
        return []
    parts = re.split(r"[,，;；、|]", line)
    return [_normalize_term(p) for p in parts if _normalize_term(p)]


def _try_parse_json_terms(text: str) -> Tuple[List[str], List[str]] | None:
    text = text.replace("```json", "").replace("```", "").strip()
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end <= start:
        return None
    try:
        obj = json.loads(text[start : end + 1])
    except json.JSONDecodeError:
        return None
    if not isinstance(obj, dict):
        return None

    def pick(keys: Sequence[str]) -> List[str]:
        for k in keys:
            v = obj.get(k)
            if isinstance(v, list):
                return [_normalize_term(str(x)) for x in v if _normalize_term(str(x))]
            if isinstance(v, str):
                return _split_list_items(v)
        return []

    terms = pick(("product_terms", "专有名词", "terms", "keywords"))
    stops = pick(("stopwords", "stop_words", "停用词", "custom_stopwords"))
    if terms or stops:
        return terms, stops
    return None


def _parse_section_lines(text: str, mode: str) -> List[str]:
    """mode: product | stop"""
    lines = text.splitlines()
    collecting = False
    items: List[str] = []
    product_hdr = re.compile(
        r"(专有名词|product[_\s-]*terms?|product\s+terms)", re.I
    )
    stop_hdr = re.compile(
        r"(专属停用|停用词|stop[_\s-]*words?|custom\s+stop)", re.I
    )
    other_hdr = re.compile(r"^#{1,3}\s+")

    for line in lines:
        raw = line.strip()
        if not raw:
            if collecting and items:
                collecting = False
            continue
        if mode == "product" and product_hdr.search(raw):
            collecting = True
            inline = re.split(r"[:：]", raw, maxsplit=1)
            if len(inline) > 1:
                items.extend(_split_list_items(inline[1]))
            continue
        if mode == "stop" and stop_hdr.search(raw):
            collecting = True
            inline = re.split(r"[:：]", raw, maxsplit=1)
            if len(inline) > 1:
                items.extend(_split_list_items(inline[1]))
            continue
        if collecting and (other_hdr.match(raw) or (
            mode == "product" and stop_hdr.search(raw)
        ) or (mode == "stop" and product_hdr.search(raw))):
            collecting = False
            continue
        if collecting:
            items.extend(_split_list_items(raw))
    return items


def _parse_llm_terms(text: str) -> Tuple[List[str], List[str]]:
    parsed = _try_parse_json_terms(text)
    if parsed:
        return parsed
    product = _parse_section_lines(text, "product")
    stop = _parse_section_lines(text, "stop")
    if not product and not stop:
        # 兜底：按行提取非空短语
        for line in text.splitlines():
            items = _split_list_items(line)
            if items and len(items) <= 8:
                product.extend(items)
    return _dedupe_terms(product), _dedupe_terms(stop)


def _dedupe_terms(items: Iterable[str]) -> List[str]:
    seen: Set[str] = set()
    out: List[str] = []
    for x in items:
        if x and x not in seen:
            seen.add(x)
            out.append(x)
    return out


def _product_name_tokens(product_name: str) -> Set[str]:
    return {
        t
        for t in re.findall(r"[a-z0-9]+", product_name.lower())
        if t
    }


# 产品标题拆词中可视为「不重要」、应参与停用的功能词（含 spaCy 英文停用词交集）
_PRODUCT_NAME_FILLER = frozenset(
    {
        "a", "an", "the", "and", "or", "but", "for", "nor", "so", "yet",
        "at", "by", "in", "of", "on", "to", "up", "as", "is", "it", "be",
        "with", "from", "into", "via", "per", "vs", "vs.",
    }
)


def _product_name_filler_tokens(product_name: str) -> Set[str]:
    """产品名拆词中的虚词/功能词 → 应停用。"""
    tokens = _product_name_tokens(product_name)
    return {t for t in tokens if t in EN_STOP_WORDS or t in _PRODUCT_NAME_FILLER}


def _product_name_core_tokens(product_name: str) -> Set[str]:
    """产品名中有分析价值的实词 → 不停用。"""
    fillers = _product_name_filler_tokens(product_name)
    return _product_name_tokens(product_name) - fillers


def _is_pure_number(s: str) -> bool:
    return bool(s) and s.isdigit()


def _finalize_term_lists(
    product_terms: List[str],
    stopwords: List[str],
    product_name: str,
) -> Tuple[List[str], Set[str]]:
    stop_set = set(EN_STOP_WORDS)
    stop_set.update(_dedupe_terms(stopwords))
    # 产品名虚词（for/the/a 等）纳入停用；实词成分保持可统计
    stop_set.update(_product_name_filler_tokens(product_name))
    for tok in _product_name_core_tokens(product_name):
        stop_set.discard(tok)
    pn_lower = _normalize_term(product_name)
    if pn_lower:
        stop_set.discard(pn_lower)

    cleaned_terms: List[str] = []
    for t in _dedupe_terms(product_terms):
        if t in stop_set:
            continue
        cleaned_terms.append(t)
    # 长短语优先匹配
    cleaned_terms.sort(key=lambda x: (-len(x.split()), -len(x)))
    return cleaned_terms, stop_set


def _load_spacy():
    import spacy

    try:
        return spacy.load("en_core_web_sm", disable=["ner", "parser"])
    except OSError as e:
        raise RuntimeError(
            "未安装 spaCy 英文模型，请执行: python3 -m spacy download en_core_web_sm"
        ) from e


def _phrase_pattern(phrase: str) -> re.Pattern[str]:
    parts = [re.escape(p) for p in phrase.split()]
    body = r"\s+".join(parts)
    return re.compile(rf"(?<!\w){body}(?!\w)", re.IGNORECASE)


def _merge_spans(spans: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
    if not spans:
        return []
    ordered = sorted(spans)
    merged: List[Tuple[int, int]] = [ordered[0]]
    for start, end in ordered[1:]:
        prev_s, prev_e = merged[-1]
        if start <= prev_e:
            merged[-1] = (prev_s, max(prev_e, end))
        else:
            merged.append((start, end))
    return merged


def _overlaps_span(char_start: int, char_end: int, spans: Sequence[Tuple[int, int]]) -> bool:
    for s, e in spans:
        if char_start < e and char_end > s:
            return True
    return False


def _apply_product_terms(
    lower: str,
    product_terms: List[str],
    counter: Counter[str],
    stop_set: Set[str],
) -> List[Tuple[int, int]]:
    """匹配专有名词短语：只计完整短语频次，并返回需屏蔽拆词统计的字符区间。"""
    spans: List[Tuple[int, int]] = []
    for phrase in product_terms:
        if not phrase or phrase in stop_set or _is_pure_number(phrase):
            continue
        pat = _phrase_pattern(phrase)
        hits = 0
        for m in pat.finditer(lower):
            spans.append((m.start(), m.end()))
            hits += 1
        if hits:
            counter[phrase] += hits
    return _merge_spans(spans)


def _tokenize_doc(
    nlp, text: str, protected_spans: Sequence[Tuple[int, int]]
) -> List[str]:
    doc = nlp(text)
    tokens: List[str] = []
    for tok in doc:
        if tok.is_space or tok.is_punct:
            continue
        char_start = tok.idx
        char_end = tok.idx + len(tok.text)
        if _overlaps_span(char_start, char_end, protected_spans):
            continue
        lemma = (tok.lemma_ or tok.text).lower().strip()
        if not lemma or _is_pure_number(lemma):
            continue
        if not re.search(r"[a-z]", lemma, re.I):
            continue
        tokens.append(lemma)
    return tokens


_SINGLE_EN_WORD = re.compile(r"^[a-z]+$")


def _should_lemma_normalize(word: str) -> bool:
    """仅对单个英文词做词形还原；多词短语、连字符短语保持原样。"""
    w = word.strip().lower()
    if not w or " " in w or "-" in w or "'" in w:
        return False
    return bool(_SINGLE_EN_WORD.match(w))


def _lemma_form(word: str, nlp) -> str:
    w = word.strip().lower()
    if not _should_lemma_normalize(w):
        return w
    doc = nlp(w)
    if not doc:
        return w
    tok = doc[0]
    if tok.is_space or tok.is_punct:
        return w
    lemma = (tok.lemma_ or tok.text).lower().strip()
    if not lemma or lemma == "-":
        return w
    return lemma


def _merge_word_forms(counter: Counter[str], nlp) -> Counter[str]:
    """写入 CSV 前合并单复数/时态等词形（如 dogs→dog, bought→buy）。"""
    merged: Counter[str] = Counter()
    for word, count in counter.items():
        merged[_lemma_form(word, nlp)] += count
    return merged


def _build_word_freq(
    rows: Sequence[Tuple[int, str]],
    product_terms: List[str],
    stop_set: Set[str],
) -> Counter[str]:
    nlp = _load_spacy()
    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_doc(nlp, text, protected):
            if tok in stop_set:
                continue
            counter[tok] += 1
    counter = Counter({k: v for k, v in counter.items() if not _is_pure_number(k)})
    return _merge_word_forms(counter, nlp)


def _save_word_freq(counter: Counter[str], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as f:
        w = csv.writer(f)
        w.writerow(["word", "count"])
        for word, count in counter.most_common():
            w.writerow([word, count])

def _step1_extract_terms(
    *,
    job_id: int,
    industry: str,
    product_name: str,
    samples: List[Tuple[int, str]],
    api_key: str,
) -> dict:
    system, user = _build_terms_prompt(industry, product_name, samples)
    logger.info("第 1 步：调用 %s 提取专有名词与停用词（样本 %s 条）", MODEL_NAME, len(samples))
    raw = _call_llm(system, user, api_key)
    product_raw, stop_raw = _parse_llm_terms(raw)
    product_terms, _stop_set = _finalize_term_lists(
        product_raw, stop_raw, product_name
    )
    logger.info(
        "解析得到专有名词 %s 条、专属停用词 %s 条",
        len(product_terms),
        len(stop_raw),
    )
    payload = {
        "job_id": job_id,
        "industry": industry,
        "product_name": product_name,
        "sample_size": len(samples),
        "sample_seed": SAMPLE_SEED,
        "sample_row_ids": [r for r, _ in samples],
        "model": MODEL_NAME,
        "product_terms": product_terms,
        "stopwords_custom": _dedupe_terms(stop_raw),
        "llm_raw": raw,
    }
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    TERMS_JSON.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("已保存 %s", TERMS_JSON)
    return payload


def _load_terms_json() -> dict:
    if not TERMS_JSON.is_file():
        raise FileNotFoundError(f"缺少 {TERMS_JSON}，请先运行完整流程或去掉 --skip-llm")
    return json.loads(TERMS_JSON.read_text(encoding="utf-8"))


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

    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",
        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 步：spaCy 全量分词与词频统计")
    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,
        "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 词频统计")
    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()