"""
从 voc_structured.sqlite 最新 job 展开 audience / pain_points / aspect / opinion /
aspect_opinion，调用 DashScope text-embedding-v4（256 维）写入 voc_embeddings.sqlite。

溯源：source_row 与结构化时一致（CSV 第 1 条数据行=1）；对应 merged_reviews_cleaned.csv
物理行号 = source_row + 1（第 1 行为表头），content 取自该数据行。

用法（项目根目录）::

  python3 向量化.py
  python3 向量化.py --workers 8
  python3 向量化.py --job-id 4 --csv merged_reviews_cleaned.csv

说明：DashScope text-embedding-v4 单次请求最多 10 条文本；脚本按批调用 API，
默认多线程并行多批（--workers），并非逐条请求。
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sqlite3
import struct
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple

from csv import DictReader
from openai import OpenAI

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

PROJECT_ROOT = Path(__file__).resolve().parent
STRUCTURED_DB = PROJECT_ROOT / "voc_structured.sqlite"
EMBED_DB = PROJECT_ROOT / "voc_embeddings.sqlite"
DEFAULT_CSV = PROJECT_ROOT / "merged_reviews_cleaned.csv"
DASHSCOPE_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
EMBEDDING_MODEL = "text-embedding-v4"
EMBEDDING_DIMENSIONS = 256
EMBED_BATCH_SIZE = 10  # DashScope text-embedding-v4 单请求 input 数组上限 10
EMBED_DEFAULT_WORKERS = 6  # 并行批次数（每批最多 10 条）
ENTITY_TYPES = (
    "audience",
    "pain_point",
    "aspect",
    "opinion",
    "aspect_opinion",
)


@dataclass
class EmbedTask:
    job_id: int
    extraction_id: int
    source_row: int
    entity_type: str
    entity_index: int
    embed_text: str
    audience: str
    aspect: str | None
    opinion: str | None
    category: str | None
    sentiment: str | None
    content: str


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 _latest_job_id(conn: sqlite3.Connection) -> int:
    row = conn.execute(
        "SELECT id FROM analysis_jobs ORDER BY id DESC LIMIT 1"
    ).fetchone()
    if not row:
        raise RuntimeError("analysis_jobs 为空，请先运行结构化")
    return int(row[0])


def _load_content_by_source_row(csv_path: Path) -> Dict[int, str]:
    """source_row（1 起）-> content；与结构化_server._load_csv_reviews 行号一致。"""
    with csv_path.open(encoding="utf-8-sig", newline="") as f:
        reader = DictReader(f)
        if not reader.fieldnames:
            raise ValueError(f"CSV 无表头: {csv_path}")
        headers = [h.strip() for h in reader.fieldnames]
        lower_map = {h.lower(): h for h in headers}
        col = None
        for name in ("content", "review", "评论"):
            if name.lower() in lower_map:
                col = lower_map[name.lower()]
                break
        if col is None:
            if len(headers) == 1:
                col = headers[0]
            else:
                raise ValueError(f"CSV 缺少 content 列: {headers}")
        out: Dict[int, str] = {}
        for idx, row in enumerate(reader, start=1):
            cell = row.get(col)
            t = str(cell or "").strip()
            if t:
                out[idx] = t
    return out


def _pack_embedding(vec: Sequence[float]) -> bytes:
    return struct.pack(f"{len(vec)}f", *vec)


def _init_embed_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS embedding_items (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            job_id INTEGER NOT NULL,
            extraction_id INTEGER NOT NULL,
            source_row INTEGER NOT NULL,
            entity_type TEXT NOT NULL,
            entity_index INTEGER NOT NULL DEFAULT 0,
            embed_text TEXT NOT NULL,
            audience TEXT,
            aspect TEXT,
            opinion TEXT,
            category TEXT,
            sentiment TEXT,
            content TEXT NOT NULL,
            dimensions INTEGER NOT NULL,
            embedding BLOB NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_embed_job_type
            ON embedding_items(job_id, entity_type);
        CREATE INDEX IF NOT EXISTS idx_embed_source_row
            ON embedding_items(job_id, source_row);
        CREATE INDEX IF NOT EXISTS idx_embed_extraction
            ON embedding_items(extraction_id, entity_type, entity_index);
        """
    )


def _should_skip_extraction(audience: str, pain_points: List[Any], feedback: List[Any]) -> bool:
    aud = (audience or "").strip().lower()
    return aud == "unknown" and not pain_points and not feedback


def _build_tasks(
    job_id: int,
    rows: List[sqlite3.Row],
    content_map: Dict[int, str],
) -> List[EmbedTask]:
    tasks: List[EmbedTask] = []
    for row in rows:
        ext_id = int(row["id"])
        source_row = int(row["source_row"])
        try:
            data = json.loads(row["extraction_json"])
        except json.JSONDecodeError as e:
            logger.warning("跳过无效 JSON extraction_id=%s: %s", ext_id, e)
            continue

        audience = str(data.get("audience", "unknown")).strip() or "unknown"
        pain_points = data.get("pain_points") or []
        feedback = data.get("product_feedback") or []
        if not isinstance(pain_points, list):
            pain_points = []
        if not isinstance(feedback, list):
            feedback = []

        if _should_skip_extraction(audience, pain_points, feedback):
            continue

        content = content_map.get(source_row, "")
        if not content:
            logger.warning(
                "source_row=%s 在 CSV 中无 content，仍写入向量但 content 为空",
                source_row,
            )

        aud_norm = audience.strip()
        if aud_norm.lower() != "unknown":
            tasks.append(
                EmbedTask(
                    job_id=job_id,
                    extraction_id=ext_id,
                    source_row=source_row,
                    entity_type="audience",
                    entity_index=0,
                    embed_text=aud_norm,
                    audience=aud_norm,
                    aspect=None,
                    opinion=None,
                    category=None,
                    sentiment=None,
                    content=content,
                )
            )

        for i, p in enumerate(pain_points):
            text = str(p).strip()
            if not text:
                continue
            tasks.append(
                EmbedTask(
                    job_id=job_id,
                    extraction_id=ext_id,
                    source_row=source_row,
                    entity_type="pain_point",
                    entity_index=i,
                    embed_text=text,
                    audience=aud_norm,
                    aspect=None,
                    opinion=None,
                    category=None,
                    sentiment=None,
                    content=content,
                )
            )

        for i, item in enumerate(feedback):
            if not isinstance(item, dict):
                continue
            aspect = str(item.get("aspect", "")).strip()
            opinion = str(item.get("opinion", "")).strip()
            category = str(item.get("category", "")).strip() or None
            sentiment = str(item.get("sentiment", "")).strip() or None
            if not aspect and not opinion:
                continue
            if aspect:
                tasks.append(
                    EmbedTask(
                        job_id=job_id,
                        extraction_id=ext_id,
                        source_row=source_row,
                        entity_type="aspect",
                        entity_index=i,
                        embed_text=aspect,
                        audience=aud_norm,
                        aspect=aspect,
                        opinion=opinion or None,
                        category=category,
                        sentiment=sentiment,
                        content=content,
                    )
                )
            if opinion:
                tasks.append(
                    EmbedTask(
                        job_id=job_id,
                        extraction_id=ext_id,
                        source_row=source_row,
                        entity_type="opinion",
                        entity_index=i,
                        embed_text=opinion,
                        audience=aud_norm,
                        aspect=aspect or None,
                        opinion=opinion,
                        category=category,
                        sentiment=sentiment,
                        content=content,
                    )
                )
            if aspect and opinion:
                merged = f"{aspect}, {opinion}"
                tasks.append(
                    EmbedTask(
                        job_id=job_id,
                        extraction_id=ext_id,
                        source_row=source_row,
                        entity_type="aspect_opinion",
                        entity_index=i,
                        embed_text=merged,
                        audience=aud_norm,
                        aspect=aspect,
                        opinion=opinion,
                        category=category,
                        sentiment=sentiment,
                        content=content,
                    )
                )
    return tasks


def _embed_one_api_batch(api_key: str, texts: List[str]) -> List[bytes]:
    """单次 API 调用（最多 batch_size 条文本）。"""
    client = OpenAI(api_key=api_key, base_url=DASHSCOPE_BASE_URL)
    resp = client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=texts,
        dimensions=EMBEDDING_DIMENSIONS,
    )
    if len(resp.data) != len(texts):
        raise RuntimeError(
            f"embedding 返回条数 {len(resp.data)} != 请求 {len(texts)}"
        )
    ordered = sorted(resp.data, key=lambda d: d.index)
    blobs: List[bytes] = []
    for item in ordered:
        vec = item.embedding
        if len(vec) != EMBEDDING_DIMENSIONS:
            raise RuntimeError(f"维度 {len(vec)} != 期望 {EMBEDDING_DIMENSIONS}")
        blobs.append(_pack_embedding(vec))
    return blobs


def _embed_batches(
    api_key: str,
    tasks: List[EmbedTask],
    *,
    batch_size: int = EMBED_BATCH_SIZE,
    workers: int = EMBED_DEFAULT_WORKERS,
) -> List[bytes]:
    """
    将任务切成每批最多 batch_size 条，并行请求 DashScope。
    返回与 tasks 顺序一致的 embedding BLOB 列表。
    """
    if batch_size < 1 or batch_size > 10:
        raise ValueError("batch_size 须在 1–10 之间（DashScope 单请求上限 10）")
    workers = max(1, workers)

    chunks: List[List[EmbedTask]] = [
        tasks[i : i + batch_size] for i in range(0, len(tasks), batch_size)
    ]
    n_chunks = len(chunks)
    if n_chunks == 0:
        return []

    logger.info(
        "共 %s 条文本，%s 批（每批≤%s 条），并行 workers=%s",
        len(tasks),
        n_chunks,
        batch_size,
        min(workers, n_chunks),
    )

    # 单线程：逻辑简单，便于限流环境
    if workers == 1:
        all_blobs: List[bytes] = []
        for i, chunk in enumerate(chunks, start=1):
            blobs = _embed_one_api_batch(api_key, [t.embed_text for t in chunk])
            all_blobs.extend(blobs)
            if i == n_chunks or i % 20 == 0:
                logger.info("Embedding 进度 %s/%s 批", i, n_chunks)
        return all_blobs

    results: List[List[bytes] | None] = [None] * n_chunks
    done = 0
    with ThreadPoolExecutor(max_workers=min(workers, n_chunks)) as pool:
        future_map = {
            pool.submit(
                _embed_one_api_batch,
                api_key,
                [t.embed_text for t in chunk],
            ): idx
            for idx, chunk in enumerate(chunks)
        }
        for fut in as_completed(future_map):
            idx = future_map[fut]
            results[idx] = fut.result()
            done += 1
            if done == n_chunks or done % 20 == 0:
                logger.info("Embedding 进度 %s/%s 批", done, n_chunks)

    return [blob for batch in results for blob in batch]  # type: ignore[union-attr]


def run_embed(
    *,
    job_id: int | None = None,
    csv_path: Path = DEFAULT_CSV,
    structured_db: Path = STRUCTURED_DB,
    embed_db: Path = EMBED_DB,
    batch_size: int = EMBED_BATCH_SIZE,
    workers: int = EMBED_DEFAULT_WORKERS,
    reset_db: bool = True,
) -> Dict[str, Any]:
    api_key = _resolve_api_key()
    if not api_key:
        raise RuntimeError(
            "缺少 DASHSCOPE_API_KEY 或项目根 .dashscope_key"
        )

    csv_path = csv_path.expanduser().resolve()
    if not csv_path.is_file():
        raise FileNotFoundError(csv_path)
    if not structured_db.is_file():
        raise FileNotFoundError(structured_db)

    content_map = _load_content_by_source_row(csv_path)
    sconn = sqlite3.connect(structured_db)
    sconn.row_factory = sqlite3.Row
    try:
        jid = job_id if job_id is not None else _latest_job_id(sconn)
        ext_rows = sconn.execute(
            """
            SELECT id, source_row, extraction_json
            FROM comment_extractions
            WHERE job_id = ?
            ORDER BY source_row
            """,
            (jid,),
        ).fetchall()
        if not ext_rows:
            raise RuntimeError(f"job_id={jid} 无 comment_extractions 记录")
    finally:
        sconn.close()

    tasks = _build_tasks(jid, ext_rows, content_map)
    if not tasks:
        raise RuntimeError("没有可向量化的条目（检查过滤规则与结构化结果）")

    by_type: Dict[str, int] = {}
    for t in tasks:
        by_type[t.entity_type] = by_type.get(t.entity_type, 0) + 1
    logger.info(
        "job_id=%s：%s 条评论结构化记录 -> %s 条向量任务 %s",
        jid,
        len(ext_rows),
        len(tasks),
        by_type,
    )

    blobs = _embed_batches(
        api_key,
        tasks,
        batch_size=batch_size,
        workers=workers,
    )

    econn = sqlite3.connect(embed_db)
    try:
        _init_embed_schema(econn)
        if reset_db:
            econn.execute("DELETE FROM embedding_items")
        else:
            econn.execute("DELETE FROM embedding_items WHERE job_id = ?", (jid,))
        econn.commit()
        insert_sql = """
            INSERT INTO embedding_items (
                job_id, extraction_id, source_row, entity_type, entity_index,
                embed_text, audience, aspect, opinion, category, sentiment,
                content, dimensions, embedding
            ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)
        """
        rows = [
            (
                task.job_id,
                task.extraction_id,
                task.source_row,
                task.entity_type,
                task.entity_index,
                task.embed_text,
                task.audience,
                task.aspect,
                task.opinion,
                task.category,
                task.sentiment,
                task.content,
                EMBEDDING_DIMENSIONS,
                blob,
            )
            for task, blob in zip(tasks, blobs)
        ]
        econn.executemany(insert_sql, rows)
        econn.commit()
        total = econn.execute("SELECT COUNT(*) FROM embedding_items").fetchone()[0]
    finally:
        econn.close()

    summary = {
        "job_id": jid,
        "structured_db": str(structured_db),
        "embed_db": str(embed_db),
        "csv": str(csv_path),
        "model": EMBEDDING_MODEL,
        "dimensions": EMBEDDING_DIMENSIONS,
        "extractions": len(ext_rows),
        "vectors": total,
        "by_entity_type": by_type,
    }
    logger.info("完成：%s", summary)
    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="VOC 结构化结果向量化")
    parser.add_argument("--job-id", type=int, default=None, help="默认最新 job")
    parser.add_argument("--csv", type=Path, default=DEFAULT_CSV)
    parser.add_argument("--structured-db", type=Path, default=STRUCTURED_DB)
    parser.add_argument("--embed-db", type=Path, default=EMBED_DB)
    parser.add_argument(
        "--batch-size",
        type=int,
        default=EMBED_BATCH_SIZE,
        help="每批 API 请求条数，最大 10（DashScope 限制）",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=EMBED_DEFAULT_WORKERS,
        help="并行批次数；设为 1 则串行",
    )
    args = parser.parse_args()
    summary = run_embed(
        job_id=args.job_id,
        csv_path=args.csv,
        structured_db=args.structured_db,
        embed_db=args.embed_db,
        batch_size=args.batch_size,
        workers=args.workers,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()