from __future__ import annotations

import hashlib
import os
from collections import deque
from typing import Deque
import vast as ast

from lib.core.VLogger import get_logger as _vlog
from lib.constants.config import mode as _ConfigMode
from lib.core.SymbolUtils import AnnotationContainsName, ExpandSoftKeywords


def compute_sha1(content: str) -> str:
    return hashlib.sha1(content.encode('utf-8')).hexdigest()[:16]


# get_file_dependencies 缓存：键 (filepath, src_root)，值 set[str]
# 单次编译中同一文件被 find_reachable_files 和 topological_sort 重复解析，
# 缓存避免重复 open+read+ast.parse+ast.walk（~350 次调用→~175 次）
_FILE_DEPS_CACHE: dict[tuple[str, str], set[str]] = {}

# Python 源文件 AST 缓存：键 filepath，值 (mtime, content, tree)
# 多处解析同一 .py 文件（get_file_dependencies、used_includes 收集、
# _collect_inline_symbols 等）共享 parse 结果，避免重复 ast.parse（~18s→~9s）
_FILE_AST_CACHE: dict[str, tuple[float, str, ast.Module]] = {}


def _check_annotation_for_export(annotation: ast.AST | None) -> bool:
    """递归检查类型注解中是否包含 CExport/CExtern/State（t.State = t.CExtern | t.CExport）"""
    if not annotation:
        return False
    return (AnnotationContainsName(annotation, 'CExport') or
            AnnotationContainsName(annotation, 'CExtern') or
            AnnotationContainsName(annotation, 'State'))


def sha1_file(filepath: str) -> str:
    with open(filepath, 'r', encoding='utf-8') as f:
        content: str = f.read()
    return compute_sha1(content)


def parse_python_file(filepath: str) -> tuple[str, ast.Module] | tuple[str, None]:
    """读取并解析 Python 文件，带 mtime 校验的 AST 缓存

    多处调用共享同一文件的 parse 结果，避免重复 open+read+ast.parse。
    返回 (content, tree)；文件读取或解析失败返回 (content, None)。
    """
    try:
        st: os.stat_result = os.stat(filepath)
    except OSError:
        return ('', None)
    mtime: float = st.st_mtime
    cached: tuple[float, str, ast.Module] | None = _FILE_AST_CACHE.get(filepath)
    if cached is not None and cached[0] == mtime:
        return (cached[1], cached[2])
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content: str = f.read()
    except Exception:
        return ('', None)
    try:
        # 软关键字展开：fn/fun/func/function -> def（行号不变），供 ast.parse 使用。
        # 注意 content 保持原始文本返回（compute_sha1 必须与 Phase1 一致），
        # 仅 tree 基于展开后的代码解析。
        tree: ast.Module = ast.parse(ExpandSoftKeywords(content))
    except SyntaxError:
        return (content, None)
    _FILE_AST_CACHE[filepath] = (mtime, content, tree)
    return (content, tree)


def get_file_dependencies(filepath: str, src_root: str) -> set[str]:
    """解析 Python 文件的导入依赖（带缓存）

    单次编译中 find_reachable_files_from_entries 和 topological_sort_files
    会对同一批文件重复调用，缓存避免重复 ast.parse（~8.6s→~4.3s）。
    """
    cache_key: tuple[str, str] = (filepath, src_root)
    cached: set[str] | None = _FILE_DEPS_CACHE.get(cache_key)
    if cached is not None:
        return cached
    dependencies: set[str] = set()
    content, tree = parse_python_file(filepath)
    if tree is not None:
        try:
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        module: str = alias.name
                        dependencies.add(module)
                elif isinstance(node, ast.ImportFrom):
                    if node.module:
                        module: str = node.module
                        if node.level > 0:
                            # from .module import name → 依赖 pkg.module（模块本身）
                            pkg_dir: str = os.path.dirname(filepath)
                            pkg_rel: str = os.path.relpath(pkg_dir, src_root).replace(os.sep, '.').replace('/', '.')
                            if pkg_rel == '.':
                                pkg_rel = ''
                            dep_module: str
                            if pkg_rel:
                                dep_module = f'{pkg_rel}.{module}'
                            else:
                                dep_module = module
                            dependencies.add(dep_module)
                        else:
                            if not module.startswith('_'):
                                dependencies.add(module)
                    elif node.level > 0 and node.names:
                        # from . import name → 依赖 pkg.name（子模块）
                        for alias in node.names:
                            pkg_dir: str = os.path.dirname(filepath)
                            pkg_rel: str = os.path.relpath(pkg_dir, src_root).replace(os.sep, '.').replace('/', '.')
                            if pkg_rel == '.':
                                pkg_rel = ''
                            dep_module: str = f'{pkg_rel}.{alias.name}' if pkg_rel else alias.name
                            dependencies.add(dep_module)
        except Exception as _e:
            if _ConfigMode == "strict":
                raise
            _vlog().warning(f"解析文件导入依赖失败: {_e}", "Exception")
    _FILE_DEPS_CACHE[cache_key] = dependencies
    return dependencies


def clear_file_deps_cache() -> None:
    """清除 get_file_dependencies 和 AST 缓存（跨编译或文件变更时调用）"""
    _FILE_DEPS_CACHE.clear()
    _FILE_AST_CACHE.clear()


def find_reachable_files_from_entries(src_root: str, entry_files: list[str]) -> set[str]:
    """从入口文件出发，找出所有可达的 .vp/.py 文件（通过导入关系）"""
    all_files: dict[str, str] = {}  # module_name -> filepath
    for root, dirs, files in os.walk(src_root):
        dirs[:] = [d for d in dirs if d not in ('__pycache__',)]
        for file in files:
            if file.endswith('.vp') or file.endswith('.py'):
                filepath: str = os.path.join(root, file)
                rel: str = os.path.relpath(filepath, src_root)
                module: str = os.path.splitext(rel)[0].replace(os.sep, '.').replace('/', '.')
                if filepath.endswith('__init__.vp') or filepath.endswith('__init__.py'):
                    all_files[module] = filepath
                    parent_module: str = module.rsplit('.', 1)[0] if '.' in module else module
                    if parent_module not in all_files or not (all_files[parent_module].endswith('__init__.vp') or all_files[parent_module].endswith('__init__.py')):
                        all_files[parent_module] = filepath
                elif module not in all_files:
                    all_files[module] = filepath
                top_module: str = module.split('.')[0]
                if top_module not in all_files:
                    all_files[top_module] = filepath

    # 后缀索引：pkg.sub.mod → 注册 "pkg.sub.mod" / "sub.mod" / "mod"，
    # 使 from-import 的短依赖名（dep）O(1) 定位目标文件，替代原
    # for key,val in all_files.items() 全表线性扫描（O(n²)，Phase1 全量
    # 重编译可达文件多时显著拖慢）。
    _suffix_index: dict[str, str] = {}
    for _key, _val in all_files.items():
        _parts: list[str] = _key.split('.')
        for _i in range(len(_parts)):
            _suf: str = '.'.join(_parts[_i:])
            if _suf not in _suffix_index:
                _suffix_index[_suf] = _val

    reachable: set[str] = set()
    queue: list[str] = list(entry_files)
    # BFS 队列用 deque：pop(0) 是 O(n) 出队（每次移动整表），
    # popleft 为 O(1)，可达文件数百时累计节省显著。
    queue_dq: Deque[str] = deque(entry_files)

    while queue_dq:
        current: str = queue_dq.popleft()
        if current in reachable:
            continue
        reachable.add(current)
        deps: set[str] = get_file_dependencies(current, src_root)
        for dep in deps:
            if dep in all_files and all_files[dep] not in reachable:
                queue_dq.append(all_files[dep])
            else:
                _target: str | None = _suffix_index.get(dep)
                if _target is not None and _target not in reachable:
                    queue_dq.append(_target)

    return reachable


def topological_sort_files(py_files: list[str], src_root: str) -> list[str]:
    """对 Python 文件进行拓扑排序，确保依赖模块先被处理"""
    # 构建模块名到文件路径的映射
    module_to_file: dict[str, str] = {}
    for filepath in py_files:
        rel: str = os.path.relpath(filepath, src_root)
        module: str = os.path.splitext(rel)[0].replace(os.sep, '.').replace('/', '.')
        module_to_file[module] = filepath
        # 也添加顶层模块名
        top_module: str = module.split('.')[0]
        if top_module not in module_to_file:
            module_to_file[top_module] = filepath

    # 构建依赖图
    graph: dict[str, set[str]] = {f: set() for f in py_files}
    for filepath in py_files:
        deps: set[str] = get_file_dependencies(filepath, src_root)
        for dep in deps:
            if dep in module_to_file:
                dep_file: str = module_to_file[dep]
                if dep_file != filepath:  # 避免自依赖
                    graph[filepath].add(dep_file)

    # 反向邻接表：依赖当前文件的文件列表。Kahn 出队时只需遍历 reverse[cur]，
    # 替代原"每出队一个节点就全图扫描找依赖者"的 O(V²)（文件多时显著拖慢）。
    reverse: dict[str, list[str]] = {f: [] for f in py_files}
    for f, deps in graph.items():
        for dep in deps:
            if dep in reverse:
                reverse[dep].append(f)

    # 拓扑排序 (Kahn's algorithm)
    in_degree: dict[str, int] = {f: 0 for f in py_files}
    for f, deps in graph.items():
        for dep in deps:
            if dep in in_degree:
                in_degree[f] += 1

    # 小根堆替代"每轮 queue.sort() + pop(0)"（O(n² log n)）：heappop 取最小
    # 保证同级别文件按字母序处理（确定性），单次 O(log n)。
    import heapq
    queue: list[str] = [f for f, degree in in_degree.items() if degree == 0]
    heapq.heapify(queue)
    sorted_files: list[str] = []

    while queue:
        current: str = heapq.heappop(queue)
        sorted_files.append(current)
        for f in reverse[current]:
            in_degree[f] -= 1
            if in_degree[f] == 0:
                heapq.heappush(queue, f)

    # 如果有环，添加剩余文件
    if len(sorted_files) < len(py_files):
        remaining: list[str] = [f for f in py_files if f not in sorted_files]
        sorted_files.extend(remaining)

    return sorted_files