"""代码图分析 — PageRank + 图扩展 + repo-map 生成

手写 PageRank，不依赖 networkx。
"""

from __future__ import annotations

import logging
import math

from .db import Database

logger = logging.getLogger(__name__)


# ── PageRank ──


def personalized_pagerank(
    adjacency: dict[int, dict[int, float]],
    personalization: dict[int, float] | None = None,
    damping: float = 0.85,
    max_iter: int = 50,
    tol: float = 1e-6,
) -> dict[int, float]:
    """PersonalizedPageRank（手写实现，无外部依赖）

    Args:
        adjacency: node → {neighbor → weight}（有向加权邻接表）
        personalization: seed nodes → bias（None = uniform）
        damping: 阻尼系数（0.85 = 15% 概率随机跳转）
        max_iter: 最大迭代次数
        tol: 收敛阈值（L1 范数变化 < tol 时停止）

    Returns:
        node → rank score（归一化 sum=1.0）
    """
    # 收集所有节点
    nodes: set[int] = set(adjacency.keys())
    for neighbors in adjacency.values():
        nodes.update(neighbors.keys())

    if not nodes:
        return {}

    n = len(nodes)
    node_list = sorted(nodes)
    idx = {node: i for i, node in enumerate(node_list)}

    # 初始化 personalization 向量
    if personalization:
        total = sum(personalization.values())
        p = [0.0] * n
        for node, bias in personalization.items():
            if node in idx:
                p[idx[node]] = bias / total if total > 0 else 1.0 / n
        # 补齐未被 personalization 覆盖的节点
        p_sum = sum(p)
        if p_sum < 1e-10:
            p = [1.0 / n] * n
        else:
            p = [v / p_sum for v in p]
    else:
        p = [1.0 / n] * n

    # 构建出边权重（归一化）
    out_weights: list[list[tuple[int, float]]] = [[] for _ in range(n)]
    for src, neighbors in adjacency.items():
        if src not in idx:
            continue
        total_w = sum(neighbors.values())
        if total_w > 0:
            for dst, w in neighbors.items():
                if dst in idx:
                    out_weights[idx[src]].append((idx[dst], w / total_w))

    # 迭代 PageRank
    rank = list(p)  # 初始 rank = personalization

    for iteration in range(max_iter):
        new_rank = [(1.0 - damping) * pv for pv in p]

        for i in range(n):
            if rank[i] > 0 and out_weights[i]:
                for j, w in out_weights[i]:
                    new_rank[j] += damping * rank[i] * w
            elif rank[i] > 0 and not out_weights[i]:
                # dangling node: 分配给所有节点（按 personalization 分布）
                for j in range(n):
                    new_rank[j] += damping * rank[i] * p[j]

        # 归一化
        total = sum(new_rank)
        if total > 0:
            new_rank = [v / total for v in new_rank]

        # 收敛检查
        diff = sum(abs(new_rank[i] - rank[i]) for i in range(n))
        rank = new_rank
        if diff < tol:
            logger.debug("PageRank 在 %d 次迭代后收敛 (diff=%.2e)", iteration + 1, diff)
            break

    return {node_list[i]: rank[i] for i in range(n)}


# ── 全局 PageRank 计算 ──


def compute_and_store_pagerank(db: Database, collection: str) -> int:
    """计算全局 PageRank 并更新 symbols 表

    1. db.get_file_graph(collection) → 文件级邻接表
    2. personalized_pagerank(adjacency) → 文件级 rank
    3. 文件 rank 分配到该文件的 symbols
    4. db.update_pageranks(ranks)

    Returns: 更新的 symbol 数
    """
    try:
        edges = db.get_file_graph(collection)
    except Exception as e:
        logger.warning("获取文件引用图失败: %s", e)
        return 0

    if not edges:
        return 0

    # 构建邻接表
    adjacency: dict[int, dict[int, float]] = {}
    for from_doc, to_doc, weight in edges:
        adjacency.setdefault(from_doc, {})[to_doc] = weight

    # 计算文件级 PageRank
    file_ranks = personalized_pagerank(adjacency)
    if not file_ranks:
        return 0

    # 将文件 rank 分配到 symbols
    symbol_ranks: dict[int, float] = {}
    try:
        for doc_id, doc_rank in file_ranks.items():
            symbols = db.get_symbols_by_doc(doc_id)
            if not symbols:
                continue
            # 均匀分配给该文件的所有 symbols
            per_symbol = doc_rank / len(symbols)
            for sym in symbols:
                symbol_ranks[sym["id"]] = per_symbol
    except Exception as e:
        logger.warning("分配 PageRank 到 symbols 失败: %s", e)
        return 0

    # 写入数据库
    if symbol_ranks:
        try:
            db.update_pageranks(symbol_ranks)
            logger.info("PageRank 已更新 %d 个 symbols", len(symbol_ranks))
        except Exception as e:
            logger.warning("更新 PageRank 失败: %s", e)
            return 0

    return len(symbol_ranks)


# ── 图扩展查询 ──


def expand_context(
    db: Database,
    seed_chunk_ids: list[int],
    hop: int = 2,
    min_confidence: float = 0.5,
    max_fanout: int = 20,
) -> list[dict]:
    """图扩展: 从种子 chunks 出发，沿引用关系扩展

    Args:
        db: 数据库实例
        seed_chunk_ids: 种子 chunk ID 列表
        hop: 最大扩展跳数
        min_confidence: 最低置信度过滤
        max_fanout: 每个节点最大扩展邻居数

    Returns:
        [{"chunk_id": int, "hop_distance": int, "via": str, "score": float}, ...]
    """
    if not seed_chunk_ids:
        return []

    # chunk_id → symbol_ids 映射
    visited_chunks: dict[int, int] = {}  # chunk_id → hop_distance
    for cid in seed_chunk_ids:
        visited_chunks[cid] = 0

    current_symbols: set[int] = set()
    results: list[dict] = []

    # 从 seed chunks 找到关联 symbols
    try:
        for cid in seed_chunk_ids:
            rows = db.conn.execute(
                "SELECT id FROM symbols WHERE chunk_id = ?", (cid,)
            ).fetchall()
            for (sym_id,) in rows:
                current_symbols.add(sym_id)
    except Exception as e:
        logger.debug("图扩展查询 symbols 失败: %s", e)
        return []

    # BFS 扩展
    for h in range(1, hop + 1):
        if not current_symbols:
            break

        next_symbols: set[int] = set()
        for sym_id in current_symbols:
            try:
                refs = db.get_refs_by_symbol(sym_id, direction="both")
            except Exception:
                continue

            # 按 confidence 排序，取 top-N
            refs = sorted(refs, key=lambda r: r.get("confidence", 0), reverse=True)
            refs = refs[:max_fanout]

            for ref in refs:
                conf = ref.get("confidence", 0)
                if conf < min_confidence:
                    continue

                target_sym_id = ref.get("symbol_id")
                if target_sym_id is None:
                    continue

                # 获取目标 symbol 的 chunk_id
                try:
                    target_rows = db.conn.execute(
                        "SELECT chunk_id FROM symbols WHERE id = ?", (target_sym_id,)
                    ).fetchall()
                except Exception:
                    continue

                for (target_chunk_id,) in target_rows:
                    if target_chunk_id is None:
                        continue
                    if target_chunk_id not in visited_chunks:
                        visited_chunks[target_chunk_id] = h
                        next_symbols.add(target_sym_id)
                        results.append({
                            "chunk_id": target_chunk_id,
                            "hop_distance": h,
                            "via": ref.get("name", ""),
                            "score": conf * (1.0 / (1.0 + h)),
                        })

        current_symbols = next_symbols

    return results


# ── Repo-Map 生成 ──


def generate_repo_map(
    db: Database,
    collection: str,
    budget: int = 1024,
    focus_files: list[str] | None = None,
) -> str:
    """生成 Aider-style repo-map

    Args:
        db: 数据库实例
        collection: 集合名称
        budget: token 预算（len(text)//4 估算）
        focus_files: 焦点文件路径列表（触发 PersonalizedPageRank）

    Returns:
        repo-map 文本
    """
    # 获取所有 symbols (kind="def")
    try:
        all_docs = db.conn.execute(
            "SELECT id, path FROM documents WHERE collection = ?",
            (collection,),
        ).fetchall()
    except Exception as e:
        logger.warning("获取文档列表失败: %s", e)
        return ""

    if not all_docs:
        return ""

    doc_path_map = {r[0]: r[1] for r in all_docs}

    # 如果有焦点文件，用 PersonalizedPageRank
    if focus_files:
        edges = db.get_file_graph(collection)
        adjacency: dict[int, dict[int, float]] = {}
        for from_doc, to_doc, weight in edges:
            adjacency.setdefault(from_doc, {})[to_doc] = weight

        # 构建 personalization
        focus_doc_ids = {}
        for fp in focus_files:
            for doc_id, path in doc_path_map.items():
                if path == fp or path.endswith(fp):
                    focus_doc_ids[doc_id] = 1.0
                    break

        file_ranks = personalized_pagerank(adjacency, personalization=focus_doc_ids or None)
    else:
        # 用全局 pagerank（从 symbols.pagerank 聚合）
        file_ranks = {}
        try:
            rows = db.conn.execute(
                """SELECT s.doc_id, SUM(s.pagerank) as total_rank
                   FROM symbols s
                   JOIN documents d ON s.doc_id = d.id
                   WHERE d.collection = ?
                   GROUP BY s.doc_id""",
                (collection,),
            ).fetchall()
            for doc_id, total_rank in rows:
                file_ranks[doc_id] = total_rank or 0.0
        except Exception as e:
            logger.warning("获取 PageRank 失败: %s", e)

    # 按 rank 排序文件
    sorted_docs = sorted(
        doc_path_map.items(),
        key=lambda x: file_ranks.get(x[0], 0.0),
        reverse=True,
    )

    # 批量查询全集合的跨文件调用（1 次 SQL，替代 N 次逐文件查询）
    all_callees = _batch_get_callees(db, collection, doc_path_map)

    # 构建 repo-map 文本（按 token 预算截断）
    lines: list[str] = []
    token_count = 0

    for doc_id, path in sorted_docs:
        try:
            symbols = db.get_symbols_by_doc(doc_id)
        except Exception:
            continue

        defs = [s for s in symbols if s.get("kind")]
        if not defs:
            continue

        file_header = f"{path}:"
        header_tokens = len(file_header) // 4 + 1
        if token_count + header_tokens > budget:
            break

        # 从批量结果中提取该文件的调用标注
        callees_by_func = _assign_callees_to_funcs(
            defs, all_callees.get(doc_id, []), doc_path_map,
        )

        file_lines = [file_header]
        file_token_est = header_tokens

        for sym in sorted(defs, key=lambda s: s.get("line_start", 0)):
            sig = sym.get("signature") or sym["name"]
            line_str = f"  {sym.get('line_start', '?')}│{sym.get('kind', 'def')} {sig}"

            # 标注调用关系
            targets = callees_by_func.get(sym["id"], [])
            if targets:
                annotation = " [→ " + ", ".join(targets[:3]) + "]"
                line_str += annotation

            line_tokens = len(line_str) // 4 + 1
            if token_count + file_token_est + line_tokens > budget:
                break
            file_lines.append(line_str)
            file_token_est += line_tokens

        if len(file_lines) > 1:  # 至少有一个 symbol
            lines.extend(file_lines)
            lines.append("")  # 空行分隔
            token_count += file_token_est

    return "\n".join(lines)


def _batch_get_callees(
    db: Database,
    collection: str,
    doc_path_map: dict[int, str],
) -> dict[int, list[tuple[int, str, int]]]:
    """批量查询全集合的跨文件调用（1 次 SQL）

    Returns:
        {from_doc_id: [(from_line, target_name, to_doc_id), ...]}
    """
    try:
        rows = db.conn.execute(
            """SELECT DISTINCT r.from_doc_id, r.from_line, s.name as target_name, r.to_doc_id
               FROM refs r
               JOIN symbols s ON r.to_symbol_id = s.id
               JOIN documents d ON r.from_doc_id = d.id
               WHERE d.collection = ?
                 AND r.to_doc_id IS NOT NULL
                 AND r.to_doc_id != r.from_doc_id""",
            (collection,),
        ).fetchall()
    except Exception as e:
        logger.debug("批量获取调用关系失败: %s", e)
        return {}

    result: dict[int, list[tuple[int, str, int]]] = {}
    for from_doc, from_line, target_name, to_doc in rows:
        result.setdefault(from_doc, []).append((from_line, target_name, to_doc))
    return result


def _assign_callees_to_funcs(
    defs: list[dict],
    refs_in_file: list[tuple[int, str, int]],
    doc_path_map: dict[int, str],
) -> dict[int, list[str]]:
    """将引用分配到所属函数，返回 {symbol_id: ["db.connect", ...]}

    行号匹配策略：用相邻函数的 line_start 构建范围区间，
    每个函数的范围 = [line_start, 下一个函数的 line_start)。
    """
    if not refs_in_file or not defs:
        return {}

    # 构建函数范围: [(sym, line_start, line_end_exclusive), ...]
    sorted_defs = sorted(defs, key=lambda s: s.get("line_start", 0))
    func_ranges: list[tuple[dict, int, int]] = []
    for i, sym in enumerate(sorted_defs):
        start = sym.get("line_start", 0)
        # 下一个函数的起始行作为当前函数的结束边界
        end = sorted_defs[i + 1].get("line_start", 0) if i + 1 < len(sorted_defs) else float("inf")
        func_ranges.append((sym, start, end))

    result: dict[int, list[str]] = {}
    seen: dict[int, set[tuple[str, int]]] = {}  # symbol_id → {(target_name, to_doc_id)}

    for from_line, target_name, to_doc_id in refs_in_file:
        if from_line is None or target_name is None:
            continue

        # 精确区间匹配: line_start <= from_line < next_line_start
        owner_sym = None
        for sym, start, end in func_ranges:
            if start <= from_line < end:
                owner_sym = sym
                break

        if owner_sym is None:
            continue

        sym_id = owner_sym["id"]
        if sym_id not in seen:
            seen[sym_id] = set()

        key = (target_name, to_doc_id)
        if key not in seen[sym_id]:
            seen[sym_id].add(key)
            target_path = doc_path_map.get(to_doc_id, "")
            basename = target_path.rsplit("/", 1)[-1].rsplit(".", 1)[0] if target_path else ""
            display = f"{basename}.{target_name}" if basename else target_name
            result.setdefault(sym_id, []).append(display)

    return result
