from pathlib import Path
import subprocess
import tempfile
import numpy as np
import networkx as nx

from .types import SimulationArgs, TemporalResult, SimulationResult
from .io_utils import process_results, prepare_network_file

def run_simulation(args: SimulationArgs, lb: float, mu: float = 1.0) -> SimulationResult:
    """
    Runs a simulation by calling the Fortran binary via fpm.

    Parameters
    ----------
    args : SimulationArgs
        Simulation arguments.
    lb : float
        Value of the infection rate.
    mu : float
        Value of the healing rate.

    Returns
    -------
    SimulationResult(
        network : NetworkFormat
            Original network specification.
        node_map : dict
            Mapping from original node IDs to Fortran node IDs.
        temporal : TemporalResult
            Temporal dynamics with:
            - t : np.ndarray
                Mean time per Gillespie tick.
            - rho_avg : np.ndarray
                Mean number of infected nodes over all runs.
            - rho_var : np.ndarray
                Variance of infected nodes.
            - n_samples : int
                Number of runs where infection is non-zero.
            - active_states : Optional[dict]
                Detailed active states per sample and time (if requested), formatted as
                {sample_id: {time: {"nodes": [...], "edges": [...]}}}.
        nx_graph : Union[nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph, None]
            Representation of the structure as a NetworkX graph, if generated.
    )
    """
    # If no output_dir is provided, create a temporary one
    if args.output_dir is None:
        tmp_context = tempfile.TemporaryDirectory()
        tmpdir = Path(tmp_context.name)
        cleanup_tmp = True
    else:
        tmpdir = Path(args.output_dir)
        tmpdir.mkdir(parents=True, exist_ok=True)
        cleanup_tmp = False

    try:
        # Prepare the network file for Fortran
        network_file_fortran, map_nodes, redundant = prepare_network_file(
            args.network
        )

        # If output_dir was provided, check for existing important output files
        # and either remove them (when overwrite_output is True) or raise an
        # explicit error to avoid accidental overwrites.
        if args.output_dir is not None:
            outp = Path(args.output_dir)
            # files the Fortran program usually writes
            candidates = [
                outp / "_results.dat",
                outp / "results.dat",
                outp / "_states_nodes.dat",
                outp / "states_nodes.dat",
            ]
            existing = [p for p in candidates if p.exists()]
            if existing:
                if getattr(args, "overwrite_output", False):
                    for p in existing:
                        try:
                            p.unlink()
                        except Exception:
                            # best-effort: continue attempting to remove others
                            pass
                else:
                    names = ", ".join(str(p.name) for p in existing)
                    raise FileExistsError(
                        f"Output directory '{outp}' already contains files: {names}. "
                        "Set SimulationArgs(overwrite_output=True) to allow removing them "
                        "automatically, or choose a different output_dir."
                    )

        # Optionally build a NetworkX graph representation
        nx_graph = None
        if args.build_nx_graph:
            file_path = Path(network_file_fortran)
            if file_path.exists():
                # read as multigraph to preserve parallel edges
                create_using = nx.MultiGraph() if redundant else nx.MultiDiGraph()
                MG = nx.read_edgelist(str(file_path), nodetype=int, create_using=create_using)

                print(str(file_path), MG.number_of_nodes(), MG.number_of_edges())

                # invert node map: fortran_id(int) -> original_id
                inverse_map = {int(fid): orig for orig, fid in map_nodes.items()} if map_nodes else {}

                # If map exists, relabel MG so multiplicity checks use original IDs
                if inverse_map:
                    try:
                        MG = nx.relabel_nodes(MG, inverse_map)
                    except Exception:
                        # keep MG as-is on failure
                        pass

                # Compare edge counts between the multigraph and its simple view.
                simple_view = nx.Graph(MG) if redundant else nx.DiGraph(MG)
                has_duplicates = MG.number_of_edges() != simple_view.number_of_edges()

                # Keep multigraph only if duplicates exist, otherwise use simple view
                G = MG if has_duplicates else simple_view
                nx_graph = G

        # Chooses initial_fraction or initial_number
        kind, value = args.initial_condition
        if kind == "number":
            initial_arg = f"--initial-number {value}"
        else:
            initial_arg = f"--initial-fraction {value}"

        # Builds the command string
        inner_cmd = (
            f"dynSIS_sampling "
            f"--output \"{tmpdir}/\" "
            f"--remove-files {args.remove_files} "
            f"--edges-file \"{network_file_fortran}\" "
            f"--edges-format { 'redundant' if redundant else 'strict' } "
            f"--algorithm {args.algorithm} "
            f"--sampler {args.sampler} "
            f"--tmax {args.tmax} "
            f"--use-qs {args.use_qs} "
            f"--n-samples {args.n_samples} "
            f"--time-scale {args.time_scale} "
            f"{initial_arg} "
            f"--lambda {lb} "
            f"--mu {mu} "
            f"--export-states {args.export_states} "
            f"--seed {args.seed} "
            f"--verbose {args.verbose} "
            f"--verbose-level {args.verbose_level} "
        )

        # Executes the command
        result = subprocess.run(inner_cmd, shell=True, capture_output=True, text=True)
        if result.returncode != 0:
            raise RuntimeError(
                f"Command failed with code {result.returncode}\n"
                f"stdout:\n{result.stdout}\n"
                f"stderr:\n{result.stderr}"
            )

        # Processes the results and returns numpy arrays
        times, rho_mean, rho_var, n_samples, active_states = process_results(tmpdir, args.export_states)

        # Create TemporalResult dataclass instance
        temporal = TemporalResult(
            t=times,          # mean time per Gillespie tick
            rho_avg=rho_mean,  # mean number of infected nodes (over all runs)
            rho_var=rho_var,   # variance of infected nodes
            n_samples=n_samples  # number of runs where rho != 0
        )

        # Create SimulationResult dataclass instance
        result = SimulationResult(
            network=args.network, # original network specification
            node_map=map_nodes, # mapping from original node IDs to Fortran node IDs
            temporal=temporal,
            active_states=active_states,  # full state trajectories if export_states is True
            nx_graph=nx_graph,  # NetworkX graph representation if requested
            stdout=result.stdout,
            stderr=result.stderr
        )

        # Return the structured result object
        return result

    finally:
        if cleanup_tmp:
            tmp_context.cleanup()
