from __future__ import annotations

import os
import time
from dataclasses import dataclass
from typing import Any

from fileatlas.utils import sha256_file

from . import DISPLAY_NAME, INTERNAL_NAME
from .schema import ArchitectureMap, Module
from .tools import InternalImportResolver, build_file_graph, build_module_graph, compute_metrics
from .tools import discover_modules, extract_public_interfaces, load_project_map
from .tools import (
    mark_shared_candidates,
    merge_strongly_coupled_modules,
    render_architecture_md_best_effort,
    split_large_modules,
    suggest_flows,
    summarize_modules_best_effort,
)
from .tools.project_map_loader import ProjectIndex, build_index
from .tools.renderer import save_json, save_text


@dataclass
class ModuleAtlasConfig:
    project_map_path: str = "project_map.json"
    out_json_path: str = "architecture.json"
    out_md_path: str = "architecture.md"
    llm_concurrency: int = 4
    # Evidence is stored on each aggregated module edge:
    #   edges[*].evidence = [{from_path,to_path,import_ref}, ...]
    # This cap prevents huge architecture.json when a repo has many repeated imports.
    max_evidence_per_edge: int = 20
    # Graph-driven merge knobs (strongly-coupled modules).
    merge_min_each_direction_weight: int = 3
    merge_min_total_bidirectional_weight: int = 10
    merge_min_fraction_of_out_weight: float = 0.40
    # Optional: split very large modules (e.g., huge core packages) into 2nd-level submodules.
    # Set <=0 to disable. A conservative default keeps behavior mostly unchanged, but helps
    # large mono-packages become more readable.
    split_large_module_file_count_threshold: int = 200


def _module_records_to_output(mods: list[Any], idx: ProjectIndex) -> list[Module]:
    out: list[Module] = []
    for r in mods:
        module_id = getattr(r, "module_id", "") if not isinstance(r, dict) else str(r.get("module_id") or "")
        name = getattr(r, "name", "") if not isinstance(r, dict) else str(r.get("name") or module_id)
        files = getattr(r, "files", []) if not isinstance(r, dict) else (r.get("files") or [])
        paths = getattr(r, "paths", []) if not isinstance(r, dict) else (r.get("paths") or [])
        if not isinstance(files, list):
            files = []
        if not isinstance(paths, list):
            paths = []

        total_size = 0
        for p in files:
            facts = idx.project.get("files", {}).get(p, {}).get("facts", {})  # type: ignore[union-attr]
            if isinstance(facts, dict):
                try:
                    total_size += int(facts.get("size_bytes") or 0)
                except Exception:
                    pass

        out.append(
            {
                "module_id": str(module_id),
                "name": str(name or module_id),
                "paths": [str(x) for x in paths] if paths else [],
                "files": [str(x) for x in files] if files else [],
                "purpose": "",
                "summary": "",
                "public_interfaces": [],
                "metrics": {
                    "file_count": int(len(files)),
                    "total_size_bytes": int(total_size),
                },
            }
        )
    return out


def build_architecture_map(cfg: ModuleAtlasConfig) -> tuple[ArchitectureMap, str]:
    project = load_project_map(cfg.project_map_path)
    idx = build_index(project)

    # Module discovery
    module_records, file_to_module = discover_modules(idx)
    module_records, file_to_module = split_large_modules(
        modules=module_records,
        file_to_module=file_to_module,
        idx=idx,
        file_count_threshold=int(cfg.split_large_module_file_count_threshold),
    )
    modules = _module_records_to_output(module_records, idx)

    # Import resolution + graphs
    resolver = InternalImportResolver(all_files=idx.all_files, package_dirs=idx.package_dirs)
    file_graph = build_file_graph(project, resolver=resolver, max_evidence_per_edge=cfg.max_evidence_per_edge)
    module_edges = build_module_graph(file_graph, file_to_module=file_to_module, max_evidence_per_edge=cfg.max_evidence_per_edge)

    # Graph-driven merge (strong coupling) and re-aggregate edges
    modules, module_edges, file_to_module = merge_strongly_coupled_modules(
        modules=modules,
        module_edges=module_edges,
        file_to_module=file_to_module,
        max_evidence_per_edge=cfg.max_evidence_per_edge,
        min_each_direction_weight=int(cfg.merge_min_each_direction_weight),
        min_total_bidirectional_weight=int(cfg.merge_min_total_bidirectional_weight),
        min_fraction_of_out_weight=float(cfg.merge_min_fraction_of_out_weight),
    )

    # Metrics
    metrics, fan_in, fan_out, fan_in_w, fan_out_w = compute_metrics(
        modules=modules,
        module_edges=module_edges,
        project=project,
        file_to_module=file_to_module,
    )

    # Attach per-module fanin/out metrics
    for m in modules:
        mid = str(m.get("module_id") or "")
        mm = m.get("metrics")
        if not isinstance(mm, dict):
            mm = {}
            m["metrics"] = mm
        mm["fan_in"] = int(fan_in.get(mid, 0))
        mm["fan_out"] = int(fan_out.get(mid, 0))
        mm["fan_in_weight"] = int(fan_in_w.get(mid, 0))
        mm["fan_out_weight"] = int(fan_out_w.get(mid, 0))

    # utils/common/shared identification (initial heuristic)
    mark_shared_candidates(modules, fan_in=fan_in, fan_out=fan_out)

    # Module-level Map summaries (LLM best-effort + fallback)
    summarize_modules_best_effort(modules, project=project, llm_concurrency=cfg.llm_concurrency)

    # Public interfaces (best-effort resolution back to symbols/span)
    extract_public_interfaces(modules=modules, idx=idx)

    # Suggested flows
    flows = suggest_flows(modules=modules, module_edges=module_edges, fan_in_w=fan_in_w, fan_out_w=fan_out_w, limit=10)

    # Reduce: render architecture.md
    meta: dict[str, Any] = {
        "generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        "root": idx.root,
        "agent": INTERNAL_NAME,
        "agent_display_name": DISPLAY_NAME,
        "version": "0.1.0",
        "project_map_path": os.path.abspath(cfg.project_map_path),
        "source_project_map_hash": "",
        "llm_enabled": True,
        "llm_concurrency": max(1, int(cfg.llm_concurrency)),
        "max_evidence_per_edge": int(cfg.max_evidence_per_edge),
    }
    try:
        meta["source_project_map_hash"] = sha256_file(cfg.project_map_path)
    except Exception:
        meta["source_project_map_hash"] = ""

    md = render_architecture_md_best_effort(
        meta=meta,
        modules=modules,
        module_edges=module_edges,
        metrics=metrics,
        suggested_flows=flows,
    )

    arch: ArchitectureMap = {
        "meta": meta,
        "modules": modules,
        "edges": module_edges,
        "metrics": metrics,
        "suggested_flows": flows,
    }
    return arch, md


def run_moduleatlas(cfg: ModuleAtlasConfig) -> ArchitectureMap:
    arch, md = build_architecture_map(cfg)
    save_json(arch, cfg.out_json_path)
    save_text(md, cfg.out_md_path)
    return arch
