from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from ..schema import ModuleRecord
from .project_map_loader import ProjectIndex, _norm_path, get_file_facts


_TWO_LEVEL_TOP_DIRS = {"src", "lib", "app", "apps", "packages", "package", "services", "service"}


def _split_parts(path: str) -> list[str]:
    p = _norm_path(path)
    return [x for x in p.split("/") if x]


def _dir_of(path: str) -> str:
    p = _norm_path(path)
    if "/" not in p:
        return ""
    return p.rsplit("/", 1)[0]


def _parent_dir(d: str) -> str:
    d = _norm_path(d)
    if not d or "/" not in d:
        return ""
    return d.rsplit("/", 1)[0]


def _module_id_for_path(path: str, idx: ProjectIndex) -> str:
    path = _norm_path(path)
    facts = get_file_facts(idx.project, path)
    lang = str(facts.get("language") or "")

    # Prefer python packages when __init__.py indicates packages.
    is_py = lang == "python" or path.endswith(".py")
    if is_py:
        d = _dir_of(path)
        if path.endswith("/__init__.py"):
            pkg_dir = d
        elif d in idx.package_dirs:
            pkg_dir = d
        else:
            pkg_dir = ""

        if pkg_dir:
            # Find the topmost package dir in a contiguous __init__.py chain.
            cur = pkg_dir
            while True:
                parent = _parent_dir(cur)
                if parent and parent in idx.package_dirs:
                    cur = parent
                    continue
                break
            return cur

    parts = _split_parts(path)
    if not parts:
        return "(root)"
    if len(parts) == 1:
        return "(root)"
    top = parts[0]
    if top in _TWO_LEVEL_TOP_DIRS and len(parts) >= 2:
        return f"{top}/{parts[1]}"
    return top


def discover_modules(idx: ProjectIndex) -> tuple[list[ModuleRecord], dict[str, str]]:
    """
    Initial, deterministic module discovery.
    Output is later refined by graph-based metrics (shared candidates, merges).
    """
    module_to_files: dict[str, list[str]] = {}
    file_to_module: dict[str, str] = {}

    for path in sorted(idx.all_files):
        mid = _module_id_for_path(path, idx)
        module_to_files.setdefault(mid, []).append(path)
        file_to_module[path] = mid

    modules: list[ModuleRecord] = []
    for mid in sorted(module_to_files.keys()):
        files = sorted(module_to_files[mid])
        paths = sorted({mid} if mid not in {"(root)"} else {"."})
        name = mid.split("/")[-1] if mid not in {"(root)"} else "(root)"
        modules.append(ModuleRecord(module_id=mid, name=name, files=files, paths=paths))

    return modules, file_to_module


def split_large_modules(
    *,
    modules: list[ModuleRecord],
    file_to_module: dict[str, str],
    idx: ProjectIndex,
    file_count_threshold: int = 200,
) -> tuple[list[ModuleRecord], dict[str, str]]:
    """
    Optional refinement pass to make architecture more detailed by splitting very large
    modules into 2nd-level submodules (module/<child_dir>).

    Rationale:
    - Some repos have a huge core package (e.g., `verl/`), and the default discovery
      collapses it into a single module, which is too coarse for documentation.
    - This split is deterministic and keeps top-level files inside the original module.

    Rules:
    - Only split modules whose file_count >= file_count_threshold.
    - For a file `mid/.../x.py`, map it to `mid/<first_child_dir>`.
    - For a file directly under `mid/` (no child dir), keep it under `mid`.
    """
    thr = int(file_count_threshold)
    if thr <= 0:
        return modules, file_to_module

    module_to_files: dict[str, list[str]] = {}
    for p, mid in (file_to_module or {}).items():
        module_to_files.setdefault(str(mid), []).append(_norm_path(str(p)))

    # Identify large modules
    large: set[str] = {mid for mid, files in module_to_files.items() if len(files) >= thr and mid not in {"(root)"}}
    if not large:
        return modules, file_to_module

    new_file_to_module: dict[str, str] = {}
    for p0, mid0 in (file_to_module or {}).items():
        p = _norm_path(str(p0))
        mid = str(mid0)
        if mid not in large:
            new_file_to_module[p] = mid
            continue
        prefix = mid.rstrip("/") + "/"
        if not p.startswith(prefix):
            new_file_to_module[p] = mid
            continue
        rest = p[len(prefix) :]
        if not rest or "/" not in rest:
            # file directly under module dir
            new_file_to_module[p] = mid
            continue
        child = rest.split("/", 1)[0].strip()
        if not child or child.startswith("."):
            new_file_to_module[p] = mid
            continue
        new_file_to_module[p] = f"{mid}/{child}"

    # Materialize new ModuleRecords
    new_module_to_files: dict[str, list[str]] = {}
    for p, mid in new_file_to_module.items():
        new_module_to_files.setdefault(mid, []).append(p)

    out_modules: list[ModuleRecord] = []
    for mid in sorted(new_module_to_files.keys()):
        files = sorted(new_module_to_files[mid])
        paths = sorted({mid} if mid not in {"(root)"} else {"."})
        name = mid.split("/")[-1] if mid not in {"(root)"} else "(root)"
        out_modules.append(ModuleRecord(module_id=mid, name=name, files=files, paths=paths))

    _ = idx  # reserved for future heuristics
    return out_modules, new_file_to_module


def mark_shared_candidates(
    modules: list[dict[str, Any]],
    *,
    fan_in: dict[str, int],
    fan_out: dict[str, int],
    min_fan_in: int = 3,
) -> None:
    """
    Mutates modules[*]["metrics"]["is_shared_candidate"] based on simple heuristic.
    """
    for m in modules:
        mid = str(m.get("module_id") or "")
        if not mid:
            continue
        fi = int(fan_in.get(mid, 0))
        fo = int(fan_out.get(mid, 0))
        is_shared = fi >= min_fan_in and fo <= max(1, fi // 3)
        metrics = m.get("metrics")
        if not isinstance(metrics, dict):
            metrics = {}
            m["metrics"] = metrics
        metrics["is_shared_candidate"] = bool(is_shared)


def merge_strongly_coupled_modules(
    *,
    modules: list[dict[str, Any]],
    module_edges: list[dict[str, Any]],
    file_to_module: dict[str, str],
    max_evidence_per_edge: int = 20,
    min_each_direction_weight: int = 3,
    min_total_bidirectional_weight: int = 10,
    min_fraction_of_out_weight: float = 0.40,
    max_iterations: int = 20,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, str]]:
    """
    Graph-driven merge pass.

    We merge A<- ->B when:
    - both directions exist
    - min(w_ab, w_ba) >= min_each_direction_weight
    - (w_ab + w_ba) >= min_total_bidirectional_weight
    - w_ab/out_weight(A) and w_ba/out_weight(B) are both >= min_fraction_of_out_weight

    Safeguards:
    - avoid merging shared-like modules (utils/common/shared) or high fan-in hubs
    - keep module_id stable by merging into the "dominant" module (larger file_count/size)
      and recording absorbed ids in module["merged_from"].
    """
    max_ev = max(0, int(max_evidence_per_edge))
    iters = max(1, int(max_iterations))

    def is_sharedish(mid: str, fan_in_count: int) -> bool:
        toks = set(mid.lower().split("/"))
        if toks & {"utils", "common", "shared"}:
            return True
        # High fan-in hubs tend to be shared; don't merge them.
        if fan_in_count >= 5:
            return True
        return False

    def recompute_edge_maps(edges: list[dict[str, Any]]) -> tuple[dict[tuple[str, str], dict[str, Any]], dict[str, int], dict[str, int]]:
        # returns (pair->edge, out_weight, fan_in_count)
        pair: dict[tuple[str, str], dict[str, Any]] = {}
        out_w: dict[str, int] = {}
        fan_in: dict[str, int] = {}
        for e in edges:
            a = str(e.get("from_module") or "")
            b = str(e.get("to_module") or "")
            w = int(e.get("weight") or 0)
            if not a or not b or a == b or w <= 0:
                continue
            pair[(a, b)] = e
            out_w[a] = out_w.get(a, 0) + w
            fan_in[b] = fan_in.get(b, 0) + 1
        return pair, out_w, fan_in

    def dominant(a: dict[str, Any], b: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
        # choose who survives
        def score(m: dict[str, Any]) -> tuple[int, int, str]:
            mm = m.get("metrics", {}) if isinstance(m.get("metrics", {}), dict) else {}
            fc = int(mm.get("file_count") or (len(m.get("files") or []) if isinstance(m.get("files"), list) else 0))
            sz = int(mm.get("total_size_bytes") or 0)
            mid = str(m.get("module_id") or "")
            return (fc, sz, mid)

        return (a, b) if score(a) >= score(b) else (b, a)

    def merge_into(dst: dict[str, Any], src: dict[str, Any]) -> None:
        dst_files = dst.get("files", [])
        src_files = src.get("files", [])
        if not isinstance(dst_files, list):
            dst_files = []
        if not isinstance(src_files, list):
            src_files = []
        dst_paths = dst.get("paths", [])
        src_paths = src.get("paths", [])
        if not isinstance(dst_paths, list):
            dst_paths = []
        if not isinstance(src_paths, list):
            src_paths = []

        dst["files"] = sorted(set([str(x) for x in dst_files] + [str(x) for x in src_files]))
        dst["paths"] = sorted(set([str(x) for x in dst_paths] + [str(x) for x in src_paths]))

        # record merge provenance
        mf = dst.get("merged_from", [])
        if not isinstance(mf, list):
            mf = []
        mf = [str(x) for x in mf if str(x)]
        mf.append(str(src.get("module_id") or ""))
        src_mf = src.get("merged_from", [])
        if isinstance(src_mf, list):
            mf.extend(str(x) for x in src_mf if str(x))
        dst["merged_from"] = sorted(set(mf))

        # update quick metrics
        mm = dst.get("metrics")
        if not isinstance(mm, dict):
            mm = {}
            dst["metrics"] = mm
        mm["file_count"] = int(len(dst["files"]))

    def remap_edges(edges: list[dict[str, Any]], redirect: dict[str, str]) -> list[dict[str, Any]]:
        agg: dict[tuple[str, str], dict[str, Any]] = {}

        def red(x: str) -> str:
            x = str(x or "")
            while x in redirect and redirect[x] != x:
                x = redirect[x]
            return x

        for e in edges:
            a0 = str(e.get("from_module") or "")
            b0 = str(e.get("to_module") or "")
            if not a0 or not b0:
                continue
            a = red(a0)
            b = red(b0)
            if not a or not b or a == b:
                continue
            w = int(e.get("weight") or 0)
            key = (a, b)
            cur = agg.get(key)
            if cur is None:
                ev = e.get("evidence", [])
                if not isinstance(ev, list):
                    ev = []
                agg[key] = {
                    "from_module": a,
                    "to_module": b,
                    "kind": "imports",
                    "weight": int(w),
                    "evidence": ev[:max_ev] if max_ev != 0 else [],
                }
            else:
                cur["weight"] = int(cur.get("weight") or 0) + int(w)
                if max_ev != 0:
                    cur_ev = cur.get("evidence", [])
                    if not isinstance(cur_ev, list):
                        cur_ev = []
                    add = e.get("evidence", [])
                    if not isinstance(add, list):
                        add = []
                    if len(cur_ev) < max_ev and add:
                        cur["evidence"] = cur_ev + add[: max(0, max_ev - len(cur_ev))]

        out = list(agg.values())
        out.sort(key=lambda x: (-int(x.get("weight") or 0), str(x.get("from_module") or ""), str(x.get("to_module") or "")))
        return out

    # iterative merges
    modules_by_id: dict[str, dict[str, Any]] = {str(m.get("module_id") or ""): m for m in modules if str(m.get("module_id") or "")}
    redirect: dict[str, str] = {}

    for _ in range(iters):
        # Keep modules list in sync
        modules = [modules_by_id[mid] for mid in sorted(modules_by_id.keys())]
        pair, out_w, fan_in_cnt = recompute_edge_maps(module_edges)

        best: tuple[int, str, str] | None = None
        for (a, b), eab in pair.items():
            eba = pair.get((b, a))
            if not eba:
                continue
            wa = int(eab.get("weight") or 0)
            wb = int(eba.get("weight") or 0)
            if min(wa, wb) < int(min_each_direction_weight):
                continue
            total = wa + wb
            if total < int(min_total_bidirectional_weight):
                continue
            out_a = max(1, int(out_w.get(a, 0)))
            out_b = max(1, int(out_w.get(b, 0)))
            if (wa / out_a) < float(min_fraction_of_out_weight):
                continue
            if (wb / out_b) < float(min_fraction_of_out_weight):
                continue
            if is_sharedish(a, int(fan_in_cnt.get(a, 0))) or is_sharedish(b, int(fan_in_cnt.get(b, 0))):
                continue
            # pick highest total first (deterministic tiebreakers)
            cand = (total, a, b) if a <= b else (total, b, a)
            if best is None or cand > best:
                best = cand

        if best is None:
            break

        _, a, b = best
        ma = modules_by_id.get(a)
        mb = modules_by_id.get(b)
        if not ma or not mb:
            break

        dst, src = dominant(ma, mb)
        dst_id = str(dst.get("module_id") or "")
        src_id = str(src.get("module_id") or "")
        if not dst_id or not src_id or dst_id == src_id:
            break

        merge_into(dst, src)
        # redirect src -> dst
        redirect[src_id] = dst_id
        modules_by_id.pop(src_id, None)

        # update file_to_module for absorbed files
        for p, mid in list(file_to_module.items()):
            if mid == src_id:
                file_to_module[p] = dst_id

        # remap module_edges to reflect merge
        module_edges = remap_edges(module_edges, redirect)

    # final materialize
    modules = [modules_by_id[mid] for mid in sorted(modules_by_id.keys())]
    return modules, module_edges, file_to_module
