"""
Code graph builder — constructs a dependency / relationship graph
from parsed files using NetworkX.

The graph captures:
  • Modules (files) as primary nodes
  • Classes / interfaces as child nodes
  • Edges: imports, method calls, inheritance (when detectable)
  • Clusters: groups of related modules (by directory / package)
"""

from __future__ import annotations

import logging
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Any

try:
    import networkx as nx
except ImportError:
    # Lightweight shim for environments without the real networkx package.
    from hld_generator import _networkx_shim as nx  # type: ignore[import-untyped]

from hld_generator.parsers.base import EntityType, ParsedFile

logger = logging.getLogger(__name__)

# Language families for cross-language edge filtering.
# Languages in the same family can legitimately import from each other;
# cross-family bare-name resolutions are almost always false positives
# (e.g. Python's `import logging` matching Go's `logging.go`).
_LANG_FAMILIES: dict[str, str] = {
    "python": "python",
    "javascript": "js",
    "typescript": "js",
    "go": "go",
    "java": "jvm",
    "kotlin": "jvm",
    "scala": "jvm",
    "c": "c",
    "cpp": "c",
    "rust": "rust",
    "ruby": "ruby",
    "csharp": "dotnet",
    "php": "php",
    "swift": "apple",
    "dart": "dart",
}


@dataclass
class ModuleNode:
    """A module (file) in the code graph."""

    id: str                         # relative file path
    language: str
    package: str                    # directory / package grouping
    classes: list[str] = field(default_factory=list)
    functions: list[str] = field(default_factory=list)
    endpoints: list[str] = field(default_factory=list)
    imports_raw: list[str] = field(default_factory=list)
    line_count: int = 0
    summary: str = ""


@dataclass
class CodeGraph:
    """Complete code graph with analysis results."""

    modules: dict[str, ModuleNode] = field(default_factory=dict)
    graph: nx.DiGraph = field(default_factory=nx.DiGraph)
    packages: dict[str, list[str]] = field(default_factory=lambda: defaultdict(list))
    languages: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    total_lines: int = 0

    # Analysis results (populated by build())
    entry_points: list[str] = field(default_factory=list)
    hub_modules: list[str] = field(default_factory=list)     # most connected
    leaf_modules: list[str] = field(default_factory=list)     # no dependents


class GraphBuilder:
    """Build a CodeGraph from parsed file results."""

    def build(self, parsed_files: list[ParsedFile], target_dir: Path | None = None) -> CodeGraph:
        cg = CodeGraph()

        # Phase 1 — create module nodes
        for pf in parsed_files:
            node = self._make_module_node(pf)
            cg.modules[node.id] = node
            cg.graph.add_node(node.id, **self._node_attrs(node))
            cg.packages[node.package].append(node.id)
            cg.languages[pf.language] += 1
            cg.total_lines += pf.line_count

        # Phase 2 — resolve import edges
        self._resolve_edges(cg, parsed_files, target_dir=target_dir)

        # Phase 3 — analyse the graph
        self._analyse(cg)

        logger.info(
            "Graph: %d modules, %d edges, %d packages",
            cg.graph.number_of_nodes(),
            cg.graph.number_of_edges(),
            len(cg.packages),
        )
        return cg

    # ── Phase 1 ────────────────────────────────────────────────────────────

    def _make_module_node(self, pf: ParsedFile) -> ModuleNode:
        path = PurePosixPath(pf.file_path)
        package = str(path.parent) if str(path.parent) != "." else "(root)"

        return ModuleNode(
            id=pf.file_path,
            language=pf.language,
            package=package,
            classes=[
                e.name for e in pf.entities
                if e.entity_type in (EntityType.CLASS, EntityType.STRUCT, EntityType.INTERFACE, EntityType.TRAIT)
            ],
            functions=[
                e.name for e in pf.entities
                if e.entity_type in (EntityType.FUNCTION, EntityType.METHOD)
            ],
            endpoints=[
                e.name for e in pf.entities
                if e.entity_type == EntityType.ENDPOINT
            ],
            imports_raw=[imp.module for imp in pf.imports],
            line_count=pf.line_count,
            summary=pf.raw_summary,
        )

    @staticmethod
    def _node_attrs(node: ModuleNode) -> dict[str, Any]:
        return {
            "language": node.language,
            "package": node.package,
            "classes": node.classes,
            "functions": node.functions,
            "endpoints": node.endpoints,
            "line_count": node.line_count,
        }

    # ── Phase 2 ────────────────────────────────────────────────────────────

    def _resolve_edges(self, cg: CodeGraph, parsed_files: list[ParsedFile], target_dir: Path | None = None) -> None:
        """Try to resolve each import to an internal module."""
        # Build a lookup of possible module identifiers.
        # Longer (more specific) paths take priority; short names are only
        # registered when unambiguous (no collision with another module).
        lookup: dict[str, str] = {}
        short_name_seen: dict[str, int] = {}

        for mod_id in cg.modules:
            p = PurePosixPath(mod_id)
            stem = str(p.with_suffix(""))
            lookup[stem] = mod_id
            lookup[stem.replace("/", ".")] = mod_id
            short_name_seen[p.stem] = short_name_seen.get(p.stem, 0) + 1

            if p.stem == "__init__":
                lookup[str(p.parent)] = mod_id
                lookup[str(p.parent).replace("/", ".")] = mod_id

        # Only register short names that are unambiguous
        for mod_id in cg.modules:
            p = PurePosixPath(mod_id)
            if short_name_seen.get(p.stem, 0) == 1:
                lookup[p.stem] = mod_id

        # Go package-level lookup: map directory paths to all files in them.
        # Go imports reference packages (directories), not individual files.
        pkg_lookup: dict[str, list[str]] = defaultdict(list)
        for mod_id in cg.modules:
            p = PurePosixPath(mod_id)
            if p.suffix == ".go":
                pkg_dir = str(p.parent)
                pkg_lookup[pkg_dir].append(mod_id)

        # Detect Go module prefix from go.mod (if available).
        go_module_prefix, go_base_dir = self._detect_go_module(parsed_files, target_dir=target_dir)

        for pf in parsed_files:
            src = pf.file_path
            src_family = _LANG_FAMILIES.get(pf.language, pf.language)
            for imp in pf.imports:
                target = self._resolve_import(
                    imp.module, lookup, pkg_lookup, go_module_prefix,
                    source_file=src, go_base_dir=go_base_dir,
                )
                if target and target != src:
                    # Reject cross-language resolutions — a Python `import logging`
                    # should never resolve to a Go `logging.go`, etc.
                    target_mod = cg.modules.get(target)
                    if target_mod:
                        target_family = _LANG_FAMILIES.get(target_mod.language, target_mod.language)
                        if src_family != target_family:
                            logger.debug(
                                "Skipping cross-language edge: %s (%s) -> %s (%s)",
                                src, pf.language, target, target_mod.language,
                            )
                            continue
                    cg.graph.add_edge(src, target, label="imports")

                    # Go package imports: a single import references the entire
                    # package directory, so add edges to ALL files in it.
                    if pf.language == "go" and pkg_lookup:
                        target_pkg_dir = str(PurePosixPath(target).parent)
                        if target_pkg_dir in pkg_lookup:
                            for sibling in pkg_lookup[target_pkg_dir]:
                                if sibling != target and sibling != src:
                                    cg.graph.add_edge(src, sibling, label="imports")

    @staticmethod
    def _read_go_mod_module(go_mod_path: Path) -> str:
        """Read the module name from a go.mod file."""
        try:
            for line in go_mod_path.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if line.startswith("module "):
                    return line[len("module "):].strip()
        except (OSError, UnicodeDecodeError):
            pass
        return ""

    @staticmethod
    def _detect_go_module(parsed_files: list[ParsedFile], target_dir: Path | None = None) -> tuple[str, str]:
        """Detect the Go module name, preferring go.mod if available.

        Returns (module_prefix, base_dir) where base_dir is the relative path
        from target_dir to the directory containing go.mod (empty string if at root).

        Strategy:
          1. Read go.mod from target_dir root and subdirectories
          2. Fall back to heuristic: find common import prefix that maps to dirs
        """
        go_files = [pf for pf in parsed_files if pf.language == "go"]
        if not go_files:
            return "", ""

        # Strategy 1: Read go.mod — check root first, then subdirectories
        if target_dir is not None:
            # Check root
            go_mod = target_dir / "go.mod"
            if go_mod.is_file():
                module_name = GraphBuilder._read_go_mod_module(go_mod)
                if module_name:
                    logger.info("Detected Go module from go.mod: %s", module_name)
                    return module_name, ""

            # Check subdirectories containing Go files
            go_dirs: set[str] = set()
            for pf in go_files:
                p = PurePosixPath(pf.file_path)
                parent = str(p.parent)
                while parent and parent != ".":
                    go_dirs.add(parent)
                    parent = str(PurePosixPath(parent).parent)

            for go_dir in sorted(go_dirs):
                go_mod = target_dir / go_dir / "go.mod"
                if go_mod.is_file():
                    module_name = GraphBuilder._read_go_mod_module(go_mod)
                    if module_name:
                        logger.info("Detected Go module from %s/go.mod: %s", go_dir, module_name)
                        return module_name, go_dir

        # Strategy 2: Heuristic — find common prefix from import paths
        all_imports: list[str] = []
        for pf in go_files:
            for imp in pf.imports:
                mod = imp.module.strip().strip("'\"")
                if "/" in mod and not mod.startswith("golang.org"):
                    all_imports.append(mod)

        if not all_imports:
            return "", ""

        known_dirs: set[str] = set()
        for pf in go_files:
            p = PurePosixPath(pf.file_path)
            parent = str(p.parent)
            while parent and parent != ".":
                known_dirs.add(parent)
                parent = str(PurePosixPath(parent).parent)

        prefix_counts: Counter[str] = Counter()
        base_dir_map: dict[str, str] = {}

        for imp_path in all_imports:
            parts = imp_path.split("/")
            for i in range(1, min(len(parts), 4)):
                suffix = "/".join(parts[i:])
                # Direct match
                if suffix in known_dirs:
                    prefix = "/".join(parts[:i])
                    prefix_counts[prefix] += 1
                    base_dir_map[prefix] = ""
                    break
                # Subdirectory match: known_dir ends with /suffix
                for kd in known_dirs:
                    if kd.endswith("/" + suffix):
                        prefix = "/".join(parts[:i])
                        prefix_counts[prefix] += 1
                        base_dir_map[prefix] = kd[: -(len(suffix) + 1)]
                        break
                else:
                    continue
                break

        if prefix_counts:
            best_prefix = prefix_counts.most_common(1)[0][0]
            return best_prefix, base_dir_map.get(best_prefix, "")
        return "", ""

    @staticmethod
    def _pick_representative_file(pkg_dir: str, files: list[str]) -> str:
        """Pick the most representative file from a Go package directory.

        Go convention: prefer a file whose stem matches the directory name
        (e.g. config/config.go for the 'config' package).
        """
        dir_name = PurePosixPath(pkg_dir).name
        for f in files:
            if PurePosixPath(f).stem == dir_name:
                return f
        return files[0]

    @staticmethod
    def _resolve_import(
        module_str: str,
        lookup: dict[str, str],
        pkg_lookup: dict[str, list[str]] | None = None,
        go_module_prefix: str = "",
        source_file: str = "",
        go_base_dir: str = "",
    ) -> str | None:
        """Attempt to match an import string to a known module.

        Tries strategies in order: direct → relative → dot-stripped →
        Go module → partial dotted → slash-suffix.
        """
        cleaned = module_str.strip().strip("'\"").rstrip(";")

        return (
            GraphBuilder._try_direct(cleaned, lookup)
            or GraphBuilder._try_relative(cleaned, lookup, source_file)
            or GraphBuilder._try_stripped(cleaned, lookup)
            or GraphBuilder._try_go_module(cleaned, lookup, pkg_lookup, go_module_prefix, go_base_dir)
            or GraphBuilder._try_partial_dotted(cleaned, lookup)
            or GraphBuilder._try_slash_suffix(cleaned, lookup, pkg_lookup)
        )

    @staticmethod
    def _try_direct(cleaned: str, lookup: dict[str, str]) -> str | None:
        """Direct lookup of the cleaned import string."""
        return lookup.get(cleaned)

    @staticmethod
    def _try_relative(cleaned: str, lookup: dict[str, str], source_file: str) -> str | None:
        """Resolve Python-style relative imports (e.g. '.utils', '..helpers')."""
        if not cleaned.startswith(".") or not source_file:
            return None

        dot_count = len(cleaned) - len(cleaned.lstrip("."))
        rel_module = cleaned.lstrip(".")
        src_path = PurePosixPath(source_file)
        base = src_path.parent
        for _ in range(dot_count - 1):
            base = base.parent

        if rel_module:
            candidate_path = str(base / rel_module.replace(".", "/"))
            if candidate_path in lookup:
                return lookup[candidate_path]
            candidate_dot = str(base).replace("/", ".") + "." + rel_module if str(base) != "." else rel_module
            if candidate_dot in lookup:
                return lookup[candidate_dot]
        else:
            candidate = str(base)
            if candidate in lookup:
                return lookup[candidate]
        return None

    @staticmethod
    def _try_stripped(cleaned: str, lookup: dict[str, str]) -> str | None:
        """Try without leading dots (non-context-aware fallback)."""
        stripped = cleaned.lstrip(".")
        return lookup.get(stripped)

    @staticmethod
    def _try_go_module(
        cleaned: str,
        lookup: dict[str, str],
        pkg_lookup: dict[str, list[str]] | None,
        go_module_prefix: str,
        go_base_dir: str = "",
    ) -> str | None:
        """Strip Go module prefix and resolve as package directory."""
        if not go_module_prefix or not cleaned.startswith(go_module_prefix + "/"):
            return None

        rel_path = cleaned[len(go_module_prefix) + 1:]
        # When Go code lives in a subdirectory (e.g. go_service/), prepend
        # the base directory to align import paths with actual file paths.
        adjusted = f"{go_base_dir}/{rel_path}" if go_base_dir else rel_path

        if adjusted in lookup:
            return lookup[adjusted]
        if pkg_lookup and adjusted in pkg_lookup:
            return GraphBuilder._pick_representative_file(adjusted, pkg_lookup[adjusted])
        # Also try without base_dir for backward compatibility
        if go_base_dir:
            if rel_path in lookup:
                return lookup[rel_path]
            if pkg_lookup and rel_path in pkg_lookup:
                return GraphBuilder._pick_representative_file(rel_path, pkg_lookup[rel_path])
        return None

    @staticmethod
    def _try_partial_dotted(cleaned: str, lookup: dict[str, str]) -> str | None:
        """Try partial match (e.g. 'myapp.utils.helpers' -> 'utils/helpers').

        Also tries suffix matching against lookup keys for cases where the
        import path is relative to a subdirectory (e.g. ``routes.health``
        should match ``python_api/routes/health`` even when the short name
        ``health`` is ambiguous across languages).
        """
        stripped = cleaned.lstrip(".")
        parts = stripped.split(".")
        for i in range(len(parts)):
            candidate = "/".join(parts[i:])
            if candidate in lookup:
                return lookup[candidate]
            candidate_dot = ".".join(parts[i:])
            if candidate_dot in lookup:
                return lookup[candidate_dot]
            # Suffix match: find lookup keys ending with /candidate.
            # Only accept unambiguous (single) matches.
            if "/" in candidate:
                suffix = "/" + candidate
                matches = [v for k, v in lookup.items() if k.endswith(suffix)]
                if len(matches) == 1:
                    return matches[0]
        return None

    @staticmethod
    def _try_slash_suffix(
        cleaned: str,
        lookup: dict[str, str],
        pkg_lookup: dict[str, list[str]] | None,
    ) -> str | None:
        """Try slash-based suffix stripping (for Go imports without module prefix)."""
        if "/" not in cleaned:
            return None
        slash_parts = cleaned.split("/")
        for i in range(1, len(slash_parts)):
            candidate = "/".join(slash_parts[i:])
            if candidate in lookup:
                return lookup[candidate]
            if pkg_lookup and candidate in pkg_lookup:
                return GraphBuilder._pick_representative_file(candidate, pkg_lookup[candidate])
        return None

    # ── Phase 3 ────────────────────────────────────────────────────────────

    def _analyse(self, cg: CodeGraph) -> None:
        """Identify entry points, hubs, and leaves."""
        in_deg = dict(cg.graph.in_degree())
        out_deg = dict(cg.graph.out_degree())

        # Entry points: detect using language-aware heuristics.
        # - Go/Java/C/C++/Dart: file contains a "main" function
        # - Python: file named __main__.py or contains "main" function
        # - Rust: file named main.rs or lib.rs
        # - Other: file named main.* or index.*
        cg.entry_points = [
            mid for mid, mod in cg.modules.items()
            if self._is_entry_point(mid, mod)
        ]

        # Hub modules: top 10% by total degree
        if cg.graph.number_of_nodes() > 0:
            degrees = sorted(
                cg.graph.nodes,
                key=lambda n: in_deg.get(n, 0) + out_deg.get(n, 0),
                reverse=True,
            )
            top_n = max(1, len(degrees) // 10)
            cg.hub_modules = degrees[:top_n]

        # Leaf modules: no outgoing edges, only imported by others
        cg.leaf_modules = [
            n for n in cg.graph.nodes
            if out_deg.get(n, 0) == 0 and in_deg.get(n, 0) > 0
        ]

    @staticmethod
    def _is_entry_point(mid: str, mod: ModuleNode) -> bool:
        """Detect whether a module is an entry point using language-aware heuristics."""
        stem = PurePosixPath(mid).stem
        lang = mod.language

        # Go/Java/C/C++/Dart: a file containing a top-level "main" function
        if lang in ("go", "java", "c", "cpp", "dart"):
            return "main" in mod.functions

        # Python: __main__.py, main.py, app.py, or files with a "main" function
        # (includes synthetic "main" from `if __name__ == "__main__":` detection)
        if lang == "python":
            return stem in ("__main__", "main") or "main" in mod.functions

        # Rust: main.rs, lib.rs, or files containing a main() function
        if lang == "rust":
            return stem in ("main", "lib") or "main" in mod.functions

        # JavaScript/TypeScript: index files or files with a top-level main
        if lang in ("javascript", "typescript"):
            return stem in ("index", "main", "app", "server")

        # Ruby: files that follow common entry point naming
        if lang == "ruby":
            return stem in ("main", "application", "config")

        # Fallback: files named main.* or index.*
        return stem in ("main", "index")

    # ── Serialisation helpers ──────────────────────────────────────────────

    @staticmethod
    def graph_summary(cg: CodeGraph) -> str:
        """Human-readable summary of the code graph for LLM context."""
        lines: list[str] = []
        lines.append("## Code Graph Summary")
        lines.append(f"- **Total modules**: {len(cg.modules)}")
        lines.append(f"- **Total lines**: {cg.total_lines:,}")
        lines.append(f"- **Languages**: {dict(cg.languages)}")
        lines.append(f"- **Packages/Dirs**: {len(cg.packages)}")
        lines.append(f"- **Dependency edges**: {cg.graph.number_of_edges()}")
        lines.append("")

        if cg.entry_points:
            lines.append("### Entry Points")
            for ep in cg.entry_points[:15]:
                lines.append(f"  - `{ep}`")
            lines.append("")

        if cg.hub_modules:
            lines.append("### Hub Modules (most connected)")
            for hm in cg.hub_modules[:10]:
                lines.append(f"  - `{hm}`")
            lines.append("")

        lines.append("### Packages")
        for pkg, members in sorted(cg.packages.items()):
            lines.append(f"  - **{pkg}/** ({len(members)} files)")
            for m in members[:5]:
                mod = cg.modules[m]
                classes_str = f" — classes: {', '.join(mod.classes)}" if mod.classes else ""
                lines.append(f"    - `{PurePosixPath(m).name}`{classes_str}")
            if len(members) > 5:
                lines.append(f"    - ... and {len(members) - 5} more")
        lines.append("")

        # Endpoints
        all_endpoints = [
            (mid, ep)
            for mid, mod in cg.modules.items()
            for ep in mod.endpoints
        ]
        if all_endpoints:
            lines.append("### API Endpoints")
            for mid, ep in all_endpoints[:30]:
                lines.append(f"  - `{ep}` ← `{PurePosixPath(mid).name}`")
            lines.append("")

        return "\n".join(lines)
