"""LLM Query Expansion — L1 查询扩展（Phase 3b）

用 LLM 生成同义词/相关术语，与主搜索并行执行。
复用 Observer provider chain（低延迟优先：Ollama→OpenClaw→OpenRouter→Gemini→Claude）。

失败时返回空列表，不影响主搜索。缓存 30min TTL。

调用链: unified_search.py → query_expand.py → observer provider.generate()
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import re
import time
from collections import OrderedDict
from dataclasses import dataclass, field

from .config import Config
from .db import tokenize_for_fts5
from .query import EnrichedQuery

logger = logging.getLogger(__name__)

# ── CJK 检测（复用 query.py 模式）──

_CJK_RE = re.compile(r"[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]")

# ── 动态扩展数量 ──


def _target_expansion_count(enriched: EnrichedQuery) -> int:
    """根据查询长度动态决定扩展词数量

    短查询(1-2 tokens) → 5: 语义模糊，需要更多补充
    中等(3-4 tokens)   → 3: 意图较明确，适量扩展
    长查询(5+ tokens)  → 2: 信息充分，少量补充
    """
    tokens = enriched.token_query.replace('"', "").split(" OR ")
    effective = [t for t in tokens if t.strip()]
    n = len(effective)
    if n <= 2:
        return 5
    if n <= 4:
        return 3
    return 2


# ── Prompt 模板 ──

_EN_PROMPT = """\
Given the code search query: "{query}"{context}
List {count} related technical terms (synonyms, related concepts, common abbreviations).
Do NOT generate generic design pattern names (Observer, Factory, Singleton, etc.) unless the query explicitly mentions them.
Reply as JSON array of strings only: ["term1", "term2", ...]"""

_ZH_PROMPT = """\
给定代码搜索查询: "{query}"{context}
列出 {count} 个相关技术术语（同义词、相关概念、常见缩写）。
不要生成通用设计模式名（Observer、Factory、Singleton 等），除非查询明确提到。
仅以 JSON 字符串数组回复: ["term1", "term2", ...]"""

# ── 停用词（text fallback 过滤）──

_STOPWORDS = frozenset({
    "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
    "have", "has", "had", "do", "does", "did", "will", "would", "shall",
    "should", "may", "might", "must", "can", "could", "of", "in", "to",
    "for", "with", "on", "at", "from", "by", "as", "or", "and", "not",
    "but", "if", "then", "else", "when", "up", "out", "no", "so",
    "than", "too", "very", "just", "about", "into", "through", "during",
    "before", "after", "above", "below", "between", "such", "each",
    "few", "more", "most", "other", "some", "only", "same", "that",
    "this", "these", "those", "it", "its", "here", "there", "also",
    "given", "code", "search", "query", "list", "related", "terms",
    "reply", "json", "array", "strings", "only", "technical",
})

# ── 缓存 ──

_expansion_cache: OrderedDict[str, tuple[list[str], float]] = OrderedDict()
_CACHE_TTL = 1800.0  # 30 分钟
_CACHE_MAX = 200

# ── Provider chain 优先级（本地模型 only）──
# expansion 在搜索热路径上，只用本地低延迟 provider。
# 没 Ollama 就不扩展 — 不值得花几十秒连云端 API。

_EXPANSION_PRIORITY = ["ollama"]


# ── 返回结构 ──


@dataclass
class ExpandResult:
    """LLM expansion 结果 + 可观测数据"""

    terms: list[str] = field(default_factory=list)
    triggered: bool = True          # 门控是否通过
    cache_hit: bool = False         # 缓存命中
    provider_name: str | None = None  # 成功的 provider
    latency_ms: int = 0             # LLM 调用延迟
    parsed_method: str = "none"     # json / text_fallback / none
    raw_count: int = 0              # LLM 原始返回词数（验证前）
    failure_reason: str | None = None  # 失败原因
    grounded_count: int = 0         # corpus-grounding 后存活词数
    pre_grounding_terms: list[str] = field(default_factory=list)  # grounding 前的原始扩展词


# ── 门控 ──


def needs_expansion(enriched: EnrichedQuery) -> bool:
    """Pre-search 门控 — 决定是否启动 LLM expansion

    触发条件: 短查询 / token 数 ≤3
    跳过条件: symbol_lookup / graph_query / 多短语 / 长查询压缩
    """
    # symbol_lookup / graph_query 已有精准路径
    if enriched.route != "text_search":
        return False
    # 已有 ≥2 个短语精确匹配，扩展价值低
    if enriched.phrases and len(enriched.phrases) >= 2:
        return False
    # 短查询：语义模糊，扩展价值高
    if enriched.short_query:
        return True
    # 长查询已经信息充分
    if enriched.compressed:
        return False
    # 默认: token 数 ≤3 → 触发
    tokens = enriched.token_query.replace('"', "").split(" OR ")
    effective_tokens = [t for t in tokens if t.strip()]
    return len(effective_tokens) <= 3


# ── JSON 解析 ──


def _parse_expansion_json(raw: str) -> list[str] | None:
    """从 LLM 输出提取 JSON 字符串数组 — 4 层降级

    Returns:
        字符串列表，或 None 表示解析失败
    """
    if not raw or not raw.strip():
        return None

    text = raw.strip()

    # Layer 1: 直接解析
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return [str(item) for item in parsed if item]
    except json.JSONDecodeError:
        pass

    # Layer 2: 修复尾逗号
    try:
        fixed = re.sub(r",\s*]", "]", text)
        parsed = json.loads(fixed)
        if isinstance(parsed, list):
            return [str(item) for item in parsed if item]
    except json.JSONDecodeError:
        pass

    # Layer 3: regex 提取 [...] 块
    start = text.find("[")
    end = text.rfind("]")
    if start != -1 and end != -1 and end > start:
        try:
            block = text[start:end + 1]
            block = re.sub(r",\s*]", "]", block)
            parsed = json.loads(block)
            if isinstance(parsed, list):
                return [str(item) for item in parsed if item]
        except json.JSONDecodeError:
            pass

    # Layer 4: 失败
    return None


def _extract_keywords_from_text(raw: str) -> list[str]:
    """从 LLM 纯文本输出中提取关键词（JSON 解析失败时的 fallback）

    提取引号内的词 + 英文标识符，过滤停用词。
    """
    if not raw:
        return []

    keywords: list[str] = []
    seen: set[str] = set()

    # 1. 提取引号内的词
    for m in re.finditer(r'["\']([a-zA-Z_]\w{1,49})["\']', raw):
        word = m.group(1)
        lower = word.lower()
        if lower not in _STOPWORDS and lower not in seen:
            keywords.append(word)
            seen.add(lower)

    # 2. 提取反引号内的词
    for m in re.finditer(r"`([a-zA-Z_]\w{1,49})`", raw):
        word = m.group(1)
        lower = word.lower()
        if lower not in _STOPWORDS and lower not in seen:
            keywords.append(word)
            seen.add(lower)

    # 3. 补充: 裸标识符（≥3 字符，非停用词）
    if len(keywords) < 3:
        for m in re.finditer(r"\b([a-zA-Z_]\w{2,49})\b", raw):
            word = m.group(1)
            lower = word.lower()
            if lower not in _STOPWORDS and lower not in seen:
                keywords.append(word)
                seen.add(lower)
                if len(keywords) >= 5:
                    break

    return keywords[:5]


# ── Validation ──


def validate_expansion(terms: list[str]) -> list[str]:
    """过滤不可信的扩展词 — 基础格式检查

    不要求在 known_symbols 中（同义词本身就是新词）。
    """
    result: list[str] = []
    seen: set[str] = set()

    for term in terms:
        if not isinstance(term, str):
            continue
        term = term.strip()
        if not term:
            continue
        # 长度限制
        if len(term) > 50 or len(term) < 2:
            continue
        # 纯数字
        if term.isdigit():
            continue
        # 去重（大小写不敏感）
        lower = term.lower()
        if lower in seen:
            continue
        seen.add(lower)
        result.append(term)
        if len(result) >= 5:
            break

    return result


# ── Corpus-Grounding 验证 ──


@dataclass
class GroundedTerm:
    """验证通过的扩展词 + 权重"""
    term: str
    weight: float  # log1p 归一化 [0, 1]
    match_count: int  # corpus + cognition 总命中数


_ALLOWED_FTS5_TABLES = frozenset({"chunks_fts", "facts_fts"})


def _fts5_match_count(conn, table: str, term: str) -> int:
    """查询 FTS5 表中某个词的匹配数 — 异常返回 0"""
    if table not in _ALLOWED_FTS5_TABLES:
        return 0
    try:
        row = conn.execute(
            f"SELECT COUNT(*) FROM {table} WHERE {table} MATCH ?",
            (term,),
        ).fetchone()
        return row[0] if row else 0
    except Exception as e:
        logger.debug("FTS5 match count 失败 %s.%s: %s", table, term, e)
        return 0


def _prepare_fts5_term(term: str) -> str:
    """准备 FTS5 查询词 — 短语加引号，单词走 tokenize

    短语检测逻辑与 query.py 对齐：含空格/下划线/连字符/CamelCase 视为短语。
    """
    stripped = term.strip()
    if not stripped:
        return ""
    # 含空格 → 短语（引号包裹）
    if " " in stripped:
        return f'"{stripped}"'
    # CamelCase → 短语
    if re.search(r"[a-z][A-Z]", stripped):
        return f'"{stripped}"'
    # 含下划线/连字符 → 短语
    if "_" in stripped or "-" in stripped:
        return f'"{stripped}"'
    # 单词 → CJK 分词后返回
    return tokenize_for_fts5(stripped)


def validate_expansion_grounded(
    terms: list[str],
    db: object | None = None,
    db_cog: object | None = None,
) -> list[GroundedTerm]:
    """Corpus-Grounding 验证 — 只保留在索引中有匹配的扩展词

    同时查 corpus (chunks_fts) 和 cognition (facts_fts)。
    返回按 weight 降序排列的 GroundedTerm 列表。

    Args:
        terms: validate_expansion() 输出的扩展词列表
        db: corpus Database（有 conn 属性）
        db_cog: cognition CogDatabase（有 conn 属性）
    """
    if not terms:
        return []

    # 无 DB 时跳过验证（降级：保留所有词，weight=1.0）
    has_corpus = db is not None and hasattr(db, "conn")
    has_cog = db_cog is not None and hasattr(db_cog, "conn")
    if not has_corpus and not has_cog:
        return [GroundedTerm(term=t, weight=1.0, match_count=0) for t in terms]

    results: list[GroundedTerm] = []
    for term in terms:
        fts5_term = _prepare_fts5_term(term)
        if not fts5_term:
            continue

        corpus_count = 0
        cog_count = 0

        if has_corpus:
            corpus_count = _fts5_match_count(db.conn, "chunks_fts", fts5_term)
        if has_cog:
            cog_count = _fts5_match_count(db_cog.conn, "facts_fts", fts5_term)

        total = corpus_count + cog_count
        if total > 0:
            weight = min(math.log1p(total) / math.log1p(50), 1.0)
            results.append(GroundedTerm(
                term=term,
                weight=round(weight, 4),
                match_count=total,
            ))

    # 按 weight 降序
    results.sort(key=lambda g: g.weight, reverse=True)
    return results


def validate_cooccurrence(
    grounded_terms: list[GroundedTerm],
    query: str,
    db: object | None = None,
) -> list[GroundedTerm]:
    """共现验证 — 扩展词必须与原始查询词在同一 chunk 中共现

    解决同形异义词问题（如 "watcher" 不与 "provider" 共现 → 过滤）。
    CJK 查询词通过 tokenize_for_fts5 分词后参与 AND 查询。

    只验证 corpus（cognition facts 通常较短，共现条件过严）。
    无 corpus db 时跳过验证（保持原样）。
    """
    if not grounded_terms or db is None or not hasattr(db, "conn"):
        return grounded_terms

    # 提取查询的 token（CJK 分词）
    tokenized_query = tokenize_for_fts5(query)
    query_tokens = [t.strip() for t in tokenized_query.split() if t.strip()]
    # 过滤太短的 token（单字符无意义）
    query_tokens = [t for t in query_tokens if len(t) >= 2]

    if not query_tokens:
        return grounded_terms

    results: list[GroundedTerm] = []
    for gt in grounded_terms:
        fts5_term = _prepare_fts5_term(gt.term)
        if not fts5_term:
            continue

        # 检查扩展词是否与任意查询词共现
        cooccurs = False
        for q_token in query_tokens:
            q_fts5 = tokenize_for_fts5(q_token)
            if not q_fts5.strip():
                continue
            try:
                and_query = f"{q_fts5} AND {fts5_term}"
                row = db.conn.execute(
                    "SELECT COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH ?",
                    (and_query,),
                ).fetchone()
                if row and row[0] > 0:
                    cooccurs = True
                    break
            except Exception as e:
                # FTS5 语法错误等 → 保守保留
                logger.debug("共现查询异常: %s", e)
                cooccurs = True
                break

        if cooccurs:
            results.append(gt)

    return results


# ── 缓存操作 ──


def _cache_key(query: str, collection: str | None) -> str:
    raw = f"{query}|{collection or ''}"
    return hashlib.sha256(raw.encode()).hexdigest()


def _cache_get(query: str, collection: str | None) -> list[str] | None:
    key = _cache_key(query, collection)
    entry = _expansion_cache.get(key)
    if entry is None:
        return None
    terms, expire_at = entry
    if time.monotonic() > expire_at:
        _expansion_cache.pop(key, None)
        return None
    _expansion_cache.move_to_end(key)
    return terms


def _cache_put(query: str, collection: str | None, terms: list[str]) -> None:
    key = _cache_key(query, collection)
    _expansion_cache[key] = (terms, time.monotonic() + _CACHE_TTL)
    _expansion_cache.move_to_end(key)
    while len(_expansion_cache) > _CACHE_MAX:
        _expansion_cache.popitem(last=False)


# ── Provider 创建 ──


def _create_expansion_providers(config: Config) -> list:
    """创建低延迟优先的 provider 列表（复用 Observer 注册表）"""
    from .cognitive.observer import _PROVIDER_REGISTRY

    # 低延迟优先（固定顺序，不反转用户的 observer 配置）
    priority = list(_EXPANSION_PRIORITY)

    providers: list = []
    import importlib
    for name in priority:
        entry = _PROVIDER_REGISTRY.get(name)
        if not entry:
            continue
        try:
            mod = importlib.import_module(entry[0], package="index1.cognitive.observer")
            cls = getattr(mod, entry[1])
            providers.append(cls(config))
        except Exception as e:
            logger.debug("Expansion provider %s 初始化失败: %s", name, e)

    return providers


# ── 主入口 ──


async def expand_query(
    query: str,
    enriched: EnrichedQuery,
    config: Config,
    db_cog: object | None = None,
    *,
    db: object | None = None,
    collection: str | None = None,
    known_symbols: set[str] | None = None,
) -> ExpandResult:
    """LLM 查询扩展 — 返回 ExpandResult（含扩展词 + 可观测数据）

    流程:
    1. 检查缓存（缓存 raw LLM 结果）
    2. 构造 prompt（语言自适应）
    3. 遍历 provider chain（低延迟优先）
    4. 解析 JSON → validate
    5. Corpus-Grounding（FTS5 验证，<1ms）
    6. 共现验证（扩展词与查询词同 chunk）

    缓存策略: raw LLM 结果缓存 30min，grounding 每次 use-time 执行
    （FTS5 <1ms 不值得缓存，reindex 自动失效）。
    失败时返回空 ExpandResult，不抛异常。
    """
    # 检查 config 开关
    if not getattr(config, "expansion_enabled", False):
        return ExpandResult(triggered=False)

    result = ExpandResult()

    # ── Phase 1: 获取 raw terms（缓存或 LLM）──

    cached = _cache_get(query, collection)
    if cached is not None:
        result.terms = list(cached)
        result.cache_hit = True
    else:
        # 语言自适应 prompt + 动态参数
        chars = query.replace(" ", "")
        cjk_count = len(_CJK_RE.findall(chars)) if chars else 0
        cjk_ratio = cjk_count / len(chars) if chars else 0.0

        count = _target_expansion_count(enriched)
        context = f"\nProject: {collection}" if collection else ""

        prompt_template = _ZH_PROMPT if cjk_ratio > 0.5 else _EN_PROMPT
        prompt = prompt_template.format(query=query, count=count, context=context)

        # 遍历 provider chain
        providers = _create_expansion_providers(config)
        if not providers:
            logger.debug("无可用 expansion provider")
            return ExpandResult(failure_reason="no providers available")

        for provider in providers:
            if not provider.available:
                continue
            try:
                t0 = time.monotonic()
                raw = await provider.generate(prompt)
                latency = int((time.monotonic() - t0) * 1000)

                if raw is None:
                    continue

                # 解析
                terms = _parse_expansion_json(raw)
                parsed_method = "json" if terms is not None else "text_fallback"
                if terms is None:
                    # JSON 失败 → text fallback
                    terms = _extract_keywords_from_text(raw)

                if not terms:
                    continue

                raw_count = len(terms)

                # 基础验证
                validated = validate_expansion(terms)
                if not validated:
                    continue

                # 缓存 raw LLM 结果（grounding 在 use-time 执行）
                _cache_put(query, collection, validated)
                logger.debug(
                    "Expansion(%s): %r → %s",
                    provider.name, query, validated,
                )
                result.terms = validated
                result.provider_name = provider.name
                result.latency_ms = latency
                result.parsed_method = parsed_method
                result.raw_count = raw_count
                break  # 成功，跳出 provider loop

            except Exception as e:
                logger.debug("Expansion provider %s 失败: %s", provider.name, e)
                continue

        # provider 全部失败
        if not result.terms and not result.cache_hit:
            _expansion_cache[_cache_key(query, collection)] = (
                [], time.monotonic() + 60.0,  # 失败结果只缓存 1 分钟
            )
            result.failure_reason = "all providers failed"

    # ── Phase 2: Corpus-Grounding（FTS5 <1ms，每次执行）──

    if result.terms:
        result.pre_grounding_terms = list(result.terms)
        grounded = validate_expansion_grounded(
            result.terms, db=db, db_cog=db_cog,
        )

        # Phase 3: 共现验证
        if grounded:
            grounded = validate_cooccurrence(grounded, query, db=db)

        result.terms = [g.term for g in grounded]
        result.grounded_count = len(grounded)

        if result.pre_grounding_terms and not result.terms:
            logger.debug(
                "Expansion grounding 过滤了所有词: %r → []",
                result.pre_grounding_terms,
            )

    return result
