import json
import os
from typing import Tuple

class FileProgressTracker:
    """Manage loading and storing latest completely processed file index

    This class save information required to continue processing in a file.
    The file structure will be:
    {
        [experiment]: {
            [self._file_ind_key]: int,
            [self._ending_ind_key]: int,
            [self._segment_id_key]: int
        }
    }
    """

    def __init__(self, save_path: str, experiment: str):
        self.path = save_path
        self.experiment = experiment
        self._file_ind_key = "file_ind"
        self._ending_ind_key = "ending_ind"
        self._segment_id_key = "segment_id"
        self._completed_key = "is_completed"

    def _load_file(self) -> dict:
        """Load processing info from file

        Returns:
            A dictionary having structure as descripted in the class info
        """
        data = {}
        if os.path.exists(self.path):
            with open(self.path) as f:
                data = json.load(f)

        if self.experiment not in data:
            data[self.experiment] = {
                self._file_ind_key: 0,
                self._ending_ind_key: 0,
                self._segment_id_key: -1,
                self._completed_key: False,
            }

        return data

    def _update_file(self, update_dict: dict) -> None:
        """Update specified keys in file"""

        data = self._load_file()
        data[self.experiment].update(update_dict)

        with open(self.path, "w+") as f:
            json.dump(data, f)

    def get_last_file_ind(self) -> Tuple[int, int, int]:
        """Get last file that was processed for this experiment

        Returns:
            A tuple containing file index, ending index in the file, and the segment number of the last processed file
        """
        data = self._load_file()
        return (
            data[self.experiment][self._file_ind_key],
            data[self.experiment][self._ending_ind_key],
            data[self.experiment][self._segment_id_key],
        )

    def update_last_file_ind(
        self, file_ind: int, ending_ind: int, segment_id: int
    ) -> None:
        """Update last file processed info in this experiment without changing other info in file if necessary"""

        self._update_file(
            {
                self._file_ind_key: file_ind,
                self._ending_ind_key: ending_ind,
                self._segment_id_key: segment_id,
            }
        )

    def mark_completion_status(self, completed: bool = True) -> None:
        self._update_file({self._completed_key: completed})

    def is_completed(self) -> bool:
        data = self._load_file()
        return data[self.experiment].get(self._completed_key, False)

    def reset_process(self) -> None:
        """Reset file processing status"""
        self.mark_completion_status(completed=False)
        self.update_last_file_ind(0, 0, -1)