from __future__ import annotations

import json
import warnings
from pathlib import Path
from typing import TYPE_CHECKING

import networkx as nx
import numpy as np

from funtracks.features import FeatureDict

if TYPE_CHECKING:
    from ..data_model import SolutionTracks, Tracks

GRAPH_FILE = "graph.json"
SEG_FILE = "seg.npy"
ATTRS_FILE = "attrs.json"


def save_tracks(tracks: Tracks, directory: Path) -> None:
    """Save the tracks to the given directory.
    Currently, saves the graph as a json file in networkx node link data format,
    saves the segmentation as a numpy npz file, and saves the time and position
    attributes and scale information in an attributes json file.
    Will make the directory if it doesn't exist.

    Args:
        tracks (Tracks): the tracks to save
        directory (Path): The directory to save the tracks in.
    """
    directory.mkdir(exist_ok=True, parents=True)
    _save_graph(tracks, directory)
    _save_seg(tracks, directory)
    _save_attrs(tracks, directory)


def _save_graph(tracks: Tracks, directory: Path) -> None:
    """Save the graph to file. Currently uses networkx node link data
    format (and saves it as json).

    Args:
        tracks (Tracks): the tracks to save the graph of
        directory (Path): The directory in which to save the graph file.
    """
    graph_file = directory / GRAPH_FILE
    graph_data = nx.node_link_data(tracks.graph, edges="links")

    def convert_np_types(data):
        """Recursively convert numpy types to native Python types."""

        if isinstance(data, dict):
            return {key: convert_np_types(value) for key, value in data.items()}
        elif isinstance(data, list):
            return [convert_np_types(item) for item in data]
        elif isinstance(data, np.ndarray):
            return data.tolist()  # Convert numpy arrays to Python lists
        elif isinstance(data, np.integer):
            return int(data)  # Convert numpy integers to Python int
        elif isinstance(data, np.floating):
            return float(data)  # Convert numpy floats to Python float
        else:
            return data  # Return the data as-is if it's already a native Python type

    graph_data = convert_np_types(graph_data)
    with open(graph_file, "w") as f:
        json.dump(graph_data, f)


def _save_seg(tracks: Tracks, directory: Path) -> None:
    """Save a segmentation as a numpy array using np.save. In the future,
    could be changed to use zarr or other file types.

    Args:
        tracks (Tracks): the tracks to save the segmentation of
        directory (Path): The directory in which to save the segmentation
    """
    if tracks.segmentation is not None:
        out_path = directory / SEG_FILE
        np.save(out_path, tracks.segmentation)


def _save_attrs(tracks: Tracks, directory: Path) -> None:
    """Save the and scale, ndim, and features in a json file in the given directory.

    Args:
        tracks (Tracks): the tracks to save the attributes of
        directory (Path):  The directory in which to save the attributes
    """
    out_path = directory / ATTRS_FILE
    attrs_dict = {
        "scale": tracks.scale
        if not isinstance(tracks.scale, np.ndarray)
        else tracks.scale.tolist(),
        "ndim": tracks.ndim,
        "features": tracks.features.dump_json(),
    }
    with open(out_path, "w") as f:
        json.dump(attrs_dict, f)


def load_tracks(
    directory: Path, seg_required: bool = False, solution: bool = False
) -> Tracks | SolutionTracks:
    """Load a Tracks object from the given directory. Looks for files
    in the format generated by Tracks.save.

    Args:
        directory (Path): The directory containing tracks to load
        seg_required (bool, optional): If true, raises a FileNotFoundError if the
            segmentation file is not present in the directory. Defaults to False.
        solution (bool, optional): If true, returns a SolutionTracks object, otherwise
            returns a normal Tracks object. Defaults to False.

    Returns:
        Tracks: A tracks object loaded from the given directory
    """
    graph_file = directory / GRAPH_FILE
    graph = _load_graph(graph_file)

    seg_file = directory / SEG_FILE
    seg = _load_seg(seg_file, seg_required=seg_required)

    attrs_file = directory / ATTRS_FILE
    attrs = _load_attrs(attrs_file)

    # filtering the warnings because the default values of time_attr and pos_attr are
    # not None. Therefore, new style Tracks attrs that have features instead of
    # pos_attr and time_attr will always trigger the warning. Updating default values
    # is breaking, and manually setting the attrs to None if features is present will
    # break if the attrs are changed/removed in the future. Can remove in v2.0.
    # Import at runtime to avoid circular dependency
    from ..data_model import SolutionTracks, Tracks

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore", message="Provided both FeatureDict and pos_attr or time_attr"
        )
        tracks: Tracks
        if solution:
            tracks = SolutionTracks(graph, seg, **attrs)
        else:
            tracks = Tracks(graph, seg, **attrs)
    return tracks


def _load_graph(graph_file: Path) -> nx.DiGraph:
    """Load the graph from the given json file. Expects networkx node_link_graph
    formatted json.

    Args:
        graph_file (Path): The json file to load into a networkx graph

    Raises:
        FileNotFoundError: If the file does not exist

    Returns:
        nx.DiGraph: A networkx graph loaded from the file.
    """
    if graph_file.is_file():
        with open(graph_file) as f:
            json_graph = json.load(f)
        return nx.node_link_graph(json_graph, directed=True, edges="links")
    else:
        raise FileNotFoundError(f"No graph at {graph_file}")


def _load_seg(seg_file: Path, seg_required: bool = False) -> np.ndarray | None:
    """Load a segmentation from a file. If the file doesn't exist, either return
    None or raise a FileNotFoundError depending on the seg_required flag.

    Args:
        seg_file (Path): The npz file to load.
        seg_required (bool, optional): If true, raise a FileNotFoundError if the
            segmentation is not present. Defaults to False.

    Returns:
        np.ndarray | None: The segmentation array, or None if it wasn't present and
            seg_required was False.
    """
    if seg_file.is_file():
        return np.load(seg_file)
    elif seg_required:
        raise FileNotFoundError(f"No segmentation at {seg_file}")
    else:
        return None


def _load_attrs(attrs_file: Path) -> dict:
    if attrs_file.is_file():
        with open(attrs_file) as f:
            json_dict = json.load(f)
        if "features" in json_dict:
            json_dict["features"] = FeatureDict.from_json(json_dict["features"])
        return json_dict

    else:
        raise FileNotFoundError(f"No attributes at {attrs_file}")


def delete_tracks(directory: Path | str):
    """Delete the tracks in the given directory. Also deletes the directory.

    Args:
        directory (Path | str): Path to a directory containing tracks
    """
    if not isinstance(directory, Path):
        directory = Path(directory)
    (directory / GRAPH_FILE).unlink()
    (directory / SEG_FILE).unlink(missing_ok=True)
    (directory / ATTRS_FILE).unlink()
    directory.rmdir()
