from pathlib import Path
from typing import Tuple, Dict, Any, List, Union, Optional
import tempfile
import numpy as np
import csv
import networkx as nx

import gzip
import bz2
import lzma

from .types import NetworkFileResult, NetworkFormat

def process_results(directory: str, export_states: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Optional[dict]]:
    """
    Reads output files generated by the Fortran binary
    and returns numpy arrays.

    Adapt this to match the actual output format.
    """
    data_file = Path(directory) / Path("_results.dat")
    data = np.loadtxt(data_file)

    times = data[:, 0]
    rho_mean = data[:, 1]
    rho_var = data[:, 2]
    n_samples = data[:, 3]

    if export_states:
        states = {}
        # the files _states_edges.dat and _states_nodes.dat should exist
        nodes_file = Path(directory) / Path("_states_nodes.dat")

        if nodes_file.exists():
            with nodes_file.open("r", encoding="utf-8") as f_nodes:
                current_sample = -1
                # Iniciamos com um valor impossível para forçar a criação da primeira amostra
                last_seen_time = -float('inf')

                for line_nodes in f_nodes:
                    parts_nodes = line_nodes.strip().split()
                    if len(parts_nodes) == 0:
                        continue

                    time = float(parts_nodes[0])
                    active_nodes = list(map(int, parts_nodes[1:]))

                    # Se o tempo atual é menor que o anterior, a amostra resetou
                    if time < last_seen_time:
                        current_sample += 1
                        states[current_sample] = {}

                    # Caso especial para a PRIMEIRA linha absoluta do arquivo
                    if current_sample == -1:
                        current_sample = 0
                        states[current_sample] = {}

                    states[current_sample][time] = {
                        "nodes": active_nodes
                    }

                    # ESSENCIAL: Atualiza o rastreador para a próxima linha
                    last_seen_time = time

        return times, rho_mean, rho_var, n_samples, states

    return times, rho_mean, rho_var, n_samples, None

def prepare_network_file(spec: NetworkFormat) -> NetworkFileResult:
    """
    Prepares the network file for the Fortran simulation.

    Parameters
    ----------
    spec : NetworkFormat
        Tuple describing the network and its format.

    Returns
    -------
    NetworkFileResult
        Paths for Fortran-ready network file, map file, and already_exists flag.
    """
    kind = spec[0]

    if kind == "edgelist":
        _, file, *rest = spec
        redundant  = rest[0] if len(rest) > 0 else False
        delimiter = rest[1] if len(rest) > 1 else None
        comment   = rest[2] if len(rest) > 2 else "#"
        cache     = rest[3] if len(rest) > 3 else False

        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"Network file not found: {file}")
        return prepare_edgelist(file, redundant, delimiter, comment, cache)

    elif kind == "fortran-edgelist":
        _, file, *rest = spec
        redundant = rest[0] if len(rest) > 0 else False

        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"Network file not found: {file}")
        return prepare_fortran_edgelist(file, redundant)

    elif kind == "networkx":
        _, source, *rest = spec
        cache = rest[0] if len(rest) > 0 else False
        if cache: print("Warning: caching is not implemented for networkx format. Ignoring cache parameter.")
        return prepare_nx(source, cache)

    elif kind == "PL":
        raise ValueError("Warning: PL format is not fully supported yet.")
        # _, gamma, N, *rest = spec
        # sample = rest[0] if len(rest) > 0 else 1
        # if gamma not in [2.5, 3.0, 6.0]:
        #     raise ValueError(f"Unsupported gamma value: {gamma}. Supported values are 2.5, 3.0, and 6.0.")
        # if N not in [100, 1000, 10000, 100000, 1000000]:
        #     raise ValueError(f"Unsupported N value: {N}. Supported values are 100, 1000, 10000, 100000, and 1000000.")
        # if sample not in [1, 2, 3, 4, 5]:
        #     raise ValueError(f"Unsupported sample value: {sample}. Supported values are 1, 2, 3, 4, and 5.")
        # base_url = 'https://zenodo.org/records/17187745/files/'
        # file_name = f"network_PL_N{N}_gamma{gamma}_sample{sample}.edgelist"
        # url = base_url + file_name
        # import requests
        # response = requests.get(url)
        # if response.status_code != 200:
        #     raise ValueError(f"Failed to download the PL network file from {url}")
        # temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".edgelist")
        # with open(temp_file.name, 'wb') as f:
        #     f.write(response.content)
        #return prepare_fortran_edgelist(temp_file.name)

    else:
        raise ValueError(f"Unknown network format: {kind}")

def prepare_edgelist(file: str, redundant: bool, delimiter: str, comment: str, cache: bool) -> NetworkFileResult:
    """
    Prepares an edgelist network file for Fortran.

    Parameters
    ----------
    file : str
        Path to the input edgelist file.
    redundant : bool
        Whether the edgelist is redundant (undirected) or not (directed).
    delimiter : str
        Delimiter used in the edgelist file.
    comment : str
        Comment character for the edgelist file.
    cache : bool
        Whether to cache the processed files.

    Returns
    -------
    file_fortran : str
        Path to the network file ready for Fortran.
    node_map : dict
        Mapping from original node IDs to Fortran node IDs.
    """
    file_fortran, map_file, already_exists = prepare_output_files(file, cache=cache)

    if already_exists:
        node_map = read_map(map_file)
        return str(file_fortran), node_map, redundant

    n_nodes = 0
    node_map = {}
    edges_mapped = []

    # read file and collect edges and node_map
    with smart_open(file, "rt") as f:
        for line in f:
            # skip comments
            if comment is not None and line.startswith(comment):
                continue

            # split line into nodes (strings)
            edge = line.strip().split(delimiter)
            mapped_edge = []
            for _n in edge:
                n = _n.strip()
                if n == "":
                    continue
                if n not in node_map:
                    n_nodes += 1
                    node_map[n] = n_nodes
                mapped_edge.append(node_map[n])
            edges_mapped.append(mapped_edge)
            if len(mapped_edge) != 2:
                print("Invalid edge:", line.strip())
                raise ValueError("Each edge must connect exactly two nodes.")

    # write edges
    write_edges(edges_mapped, file_fortran)

    # write node_map
    write_map(node_map, map_file)

    return str(file_fortran), node_map, redundant

def prepare_fortran_edgelist(file: str, redundant: bool) -> NetworkFileResult:
    # we do not create a new file!
    return str(file), {}, redundant

def prepare_nx(graph: Union[nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph], cache: bool) -> NetworkFileResult:
    """
    Prepare a NetworkX graph for Fortran.

    This function accepts both undirected (Graph/MultiGraph) and directed
    (DiGraph/MultiDiGraph) NetworkX objects. The returned `redundant` flag
    follows the same convention as other `prepare_*` functions: True means the
    input is undirected/redundant (edges appear twice), False means directed.
    """
    # use a temp file name, randomized
    temp_file_name = tempfile.mkstemp()[1]
    file_fortran, map_file, already_exists = prepare_output_files(temp_file_name, cache=cache)

    # infer whether the graph is undirected (redundant) or directed
    redundant = not graph.is_directed()

    if already_exists:
        node_map = read_map(map_file)
        return str(file_fortran), node_map, redundant

    # If not cached, we need to process the graph
    node_map = {}
    edges_mapped = []
    n_nodes = 0

    # iterate over edges; for multigraphs this will include parallel edges
    for ori_edge in graph.edges:
        # supress third collumn for multigraphs
        edge = (ori_edge[0], ori_edge[1])
        mapped_edge = []
        for n in edge:
            if n not in node_map:
                n_nodes += 1
                node_map[n] = n_nodes
            mapped_edge.append(node_map[n])
        edges_mapped.append(mapped_edge)
        if len(mapped_edge) != 2:
            print("Invalid edge:", edge)
            raise ValueError("Each edge must connect exactly two nodes.")

    # write edges and node map
    write_edges(edges_mapped, file_fortran)
    write_map(node_map, map_file)

    return str(file_fortran), node_map, redundant

def prepare_output_files(input_file: str, prefix: str = "", cache: bool = False) -> Tuple[Path, Path, bool]:
    """
    Prepares the paths for Fortran-ready network files.

    Parameters
    ----------
    input_file : str
        Path to the original input file.
    prefix : str
        Optional prefix for the generated files (default: '').

    Returns
    -------
    file_fortran : Path
        Path to the Fortran-ready network file.
    map_file : Path
        Path to the node map file.
    already_exists : bool
        True if both files already exist.
    """
    input_path = Path(input_file)
    dirname = input_path.parent
    basename = input_path.name

    file_fortran = dirname / Path(f"{prefix}.{basename}_edges.dynSIS")
    map_file = dirname / Path(f"{prefix}.{basename}_map_nodes.dynSIS")

    # Check if both files already exist, but is False if not cache
    if file_fortran.exists() and map_file.exists():
        return file_fortran, map_file, True and cache

    # Test if we can write to the directory
    try:
        test_file = dirname / Path("write_test")
        with open(test_file, "w") as f:
            f.write("test")
        test_file.unlink()
    except Exception:
        # Fallback to temporary directory
        tmpdir = Path(tempfile.mkdtemp())
        file_fortran = tmpdir / Path(f"{basename}_edges.dynSIS")
        map_file = tmpdir / Path(f"{basename}_map_nodes.dynSIS")

    return file_fortran, map_file, False

def write_map(mapping: Dict[Any, int], file_path: str) -> None:
    """
    Writes a mapping to a CSV file.

    Parameters
    ----------
    mapping : dict
        Dictionary mapping any hashable type to int.
    file_path : str
        Path to the CSV file to write.
    """
    file_path = Path(file_path)
    with file_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        # write header
        writer.writerow(["original_id", "fortran_id"])
        # write data
        for key, value in mapping.items():
            writer.writerow([key, value])


def read_map(file_path: str) -> Dict[str, int]:
    """
    Reads a mapping from a CSV file.

    Parameters
    ----------
    file_path : str
        Path to the CSV file to read.

    Returns
    -------
    mapping : dict
        Dictionary mapping string keys to int values.
    """
    mapping = {}
    file_path = Path(file_path)
    with file_path.open("r", newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            key = row["original_id"]
            value = int(row["fortran_id"])
            mapping[key] = value
    return mapping

def write_edges(edges: List[List[int]], file_path: str) -> None:
    """
    Writes a list of edges to a plain text file.

    Each sublist in `edges` represents an edge (or hyperedge),
    and the integers are separated by spaces. No header is written.

    Parameters
    ----------
    edges : List[List[int]]
        List of edges (or hyperedges), each as a list of node IDs.
    file_path : str
        Path to the file to write.
    """
    file_path = Path(file_path)
    with file_path.open("w", encoding="utf-8") as f:
        for edge in edges:
            line = " ".join(str(node) for node in edge)
            f.write(f"{line}\n")

def smart_open(file: str, mode: str = "rt"):
    """
    Opens a file, automatically handling compression based on extension.
    Supports .gz, .bz2, .xz (lzma) and regular files.
    """
    path = Path(file)
    if path.suffix == ".gz":
        return gzip.open(file, mode)
    elif path.suffix == ".bz2":
        return bz2.open(file, mode)
    elif path.suffix == ".xz":
        return lzma.open(file, mode)
    else:
        return open(file, mode)