from pathlib import Path
from typing import Tuple, Dict, Any, List, Union, Optional
import tempfile
import numpy as np
import csv
import xgi
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
        edges_file = Path(directory) / Path("_states_edges.dat")
        nodes_file = Path(directory) / Path("_states_nodes.dat")

        if edges_file.exists() and nodes_file.exists():
            # each file contains the first column as time, the rest the list of active edges/nodes
            # time is crescent, and when it changes, it means a new sample
            with edges_file.open("r", encoding="utf-8") as f_edges, nodes_file.open("r", encoding="utf-8") as f_nodes:
                current_sample = -1
                current_time = None
                for line_edges, line_nodes in zip(f_edges, f_nodes):
                    parts_edges = line_edges.strip().split()
                    parts_nodes = line_nodes.strip().split()
                    if len(parts_edges) == 0 or len(parts_nodes) == 0:
                        continue
                    time = float(parts_edges[0]) # same as parts_nodes[0]
                    active_edges = list(map(int, parts_edges[1:]))
                    active_nodes = list(map(int, parts_nodes[1:]))

                    if current_time is None or time < current_time:
                        # new sample
                        current_sample += 1
                        states[current_sample] = {}
                        current_time = time

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

        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
        delimiter = rest[0] if len(rest) > 0 else None
        comment   = rest[1] if len(rest) > 1 else "#"
        cache     = rest[2] if len(rest) > 2 else False

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

    elif kind == "fortran-edgelist":
        _, file = spec
        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"Network file not found: {file}")
        return prepare_fortran_edgelist(file)

    elif kind == "bipartite":
        _, file, *rest = spec
        delimiter = rest[0] if len(rest) > 0 else None
        comment   = rest[1] if len(rest) > 1 else "#"
        cache     = rest[2] if len(rest) > 2 else False
        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"Network file not found: {file}")
        return prepare_bipartite(file, delimiter, comment, cache)

    # TODO: #1 add cache creation for xgi, xgi_json, hif and networkx
    elif kind == "xgi":
        _, source, *rest = spec
        cache = rest[0] if len(rest) > 0 else False
        if cache: print("Warning: caching is not implemented for xgi format. Ignoring cache parameter.")
        return prepare_xgi(source, cache)

    elif kind == "xgi_json":
        _, file, *rest = spec
        cache = rest[0] if len(rest) > 0 else False
        if cache: print("Warning: caching is not implemented for xgi_json format. Ignoring cache parameter.")
        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"XGI JSON file not found: {file}")
        return prepare_xgi_json(file, cache)

    elif kind == "hif":
        _, file, *rest = spec
        cache = rest[0] if len(rest) > 0 else False
        if cache: print("Warning: caching is not implemented for hif format. Ignoring cache parameter.")
        file_path = Path(file)
        if not file_path.exists():
            raise FileNotFoundError(f"HIF network file not found: {file}")
        return prepare_hif(file, cache)

    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":
        _, 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, delimiter: str, comment: str, cache: bool) -> NetworkFileResult:
    """
    Prepares an edgelist network file for Fortran.

    Parameters
    ----------
    file : str
        Path to the input edgelist file.

    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

    n_nodes = 0
    node_map = {}
    edges_mapped = []
    max_num_nodes_in_an_edge = 0

    # 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) > max_num_nodes_in_an_edge:
                max_num_nodes_in_an_edge = len(mapped_edge)

    # write edges
    write_edges(edges_mapped, file_fortran)

    # write node_map
    write_map(node_map, map_file)

    if max_num_nodes_in_an_edge < 2:
        raise ValueError("The edgelist must contain at least one edge with two nodes. Please check your input file and delimiter setting.")

    return str(file_fortran), node_map

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

def prepare_bipartite(file: str, delimiter: str, comment: str, cache: bool) -> NetworkFileResult:
    file_fortran, map_file, already_exists = prepare_output_files(file, cache=cache)

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

    # If not cached, we need to process the file
    # In this case, we consider that each line contains node: edge
    node_map = {}
    edges_mapped = {}
    with smart_open(file, "rt") as f:
        for line in f:
            if comment is not None and line.startswith(comment):
                continue
            # Here we expect lines to be in the format "node: edge"
            parts = line.strip().split(delimiter)
            if len(parts) != 2: # we just ignore malformed lines
                continue
            _node, _edge = parts
            node = _node.strip()
            edge = _edge.strip()
            if node == "" or edge == "":
                continue
            if node not in node_map:
                node_map[node] = len(node_map) + 1
            edges_mapped[edge] = edges_mapped.get(edge, []) + [node_map[node]]

    edges_mapped = list(edges_mapped.values())
    max_num_nodes_in_an_edge = max(len(edge) for edge in edges_mapped) if edges_mapped else 0

    # write edges
    write_edges(edges_mapped, file_fortran)

    # write node_map
    write_map(node_map, map_file)

    if max_num_nodes_in_an_edge < 2:
        raise ValueError("The edgelist must contain at least one edge with two nodes. Please check your input file and delimiter setting.")

    return file_fortran, node_map

def prepare_xgi_object(H: xgi.core.hypergraph.Hypergraph, cache: bool) -> NetworkFileResult:
    # 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)

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

    edgelist = xgi.readwrite.edgelist.generate_edgelist(H)

    node_map = {}
    edges_mapped = []
    n_nodes = 0
    max_num_nodes_in_an_edge = 0

    for edge in edgelist:
        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) > max_num_nodes_in_an_edge:
            max_num_nodes_in_an_edge = len(mapped_edge)

    # write edges
    write_edges(edges_mapped, file_fortran)

    # write node_map
    write_map(node_map, map_file)

    if max_num_nodes_in_an_edge < 2:
        raise ValueError("The hypergraph must contain at least one edge with two nodes. Please check your input file.")

    return file_fortran, node_map

def prepare_nx(graph: nx.Graph, cache: bool) -> NetworkFileResult:
    # 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)

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

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

    for edge in graph.edges:
        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) > max_num_nodes_in_an_edge:
            max_num_nodes_in_an_edge = len(mapped_edge)

    # write edges
    write_edges(edges_mapped, file_fortran)

    # write node_map
    write_map(node_map, map_file)

    if max_num_nodes_in_an_edge < 2:
        raise ValueError("The graph must contain at least one edge with two nodes. Please check your input.")

    return file_fortran, node_map

def prepare_xgi(name_or_object: Union[str, xgi.core.hypergraph.Hypergraph], cache: bool) -> NetworkFileResult:

    # check if file is a path or a name of a pre-loaded object
    if isinstance(name_or_object, xgi.core.hypergraph.Hypergraph):
        H = name_or_object
    else:
        H = xgi.readwrite.xgi_data.load_xgi_data(name_or_object)

    return prepare_xgi_object(H, cache)

def prepare_xgi_json(file: str, cache: bool) -> NetworkFileResult:
    file_fortran, map_file, already_exists = prepare_output_files(file, cache=cache)

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

    H = xgi.readwrite.json.read_json(file)

    return prepare_xgi_object(H)


def prepare_hif(file: str, cache: bool) -> NetworkFileResult:
    file_fortran, map_file, already_exists = prepare_output_files(file, cache=cache)

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

    H = xgi.readwrite.hif.read_hif(file)

    return prepare_xgi_object(H)

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.hyperSIS")
    map_file = dirname / Path(f"{prefix}.{basename}_map_nodes.hyperSIS")

    # 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.hyperSIS")
        map_file = tmpdir / Path(f"{basename}_map_nodes.hyperSIS")

    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)