from __future__ import annotations

import os
import shutil
import zipfile
import pandas as pd

from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, List, Union
from unidecode import unidecode

from data_to_paper.code_and_output_files.file_view_params import ViewPurpose
from data_to_paper.env import FOLDER_FOR_RUN
from data_to_paper.latex.clean_latex import wrap_as_latex_code_output
from data_to_paper.utils.file_utils import run_in_directory, clear_directory
from data_to_paper.utils.mutable import Mutable
from data_to_paper.code_and_output_files.referencable_text import NumericReferenceableText, \
    hypertarget_if_referencable_text_product, ReferencableTextProduct

TEXT_EXTS = ['.txt', '.md', '.csv', '.xls', '.xlsx']


@dataclass(frozen=True)
class DataFileDescription:
    file_path: str  # relative to the data directory.  should normally just be the file name
    description: Optional[Union[str, ReferencableTextProduct]] = None  # a user provided description of the file
    originated_from: Optional[str] = None  # None for raw file
    is_binary: Optional[bool] = None  # None for auto based on file extension

    def get_is_binary(self):
        """
        Return True if the file is binary.
        """
        if self.is_binary is not None:
            return self.is_binary

        return Path(self.file_path).suffix not in TEXT_EXTS

    def is_excel(self):
        return Path(self.file_path).suffix in ['.xlsx', '.xls']

    def get_file_header(self, num_lines: int = 4):
        """
        Return the first `num_lines` lines of the file (if they exist).
        """
        if self.is_excel():
            # go over all sheets and return all of them:
            df = pd.read_excel(self.file_path, sheet_name=None)
            s = f'This is an Excel file with {len(df)} sheets. Here is the first few rows for each sheet:\n\n'
            for sheet_name in df.keys():
                s += f'### Sheet: "{sheet_name}"\n'
                s += f'```output\n{df[sheet_name].head(num_lines).to_string(index=False)}\n```\n'
            s += '\n'
            return s

        with open(self.file_path) as f:
            head = []
            for _ in range(num_lines):
                try:
                    head.append(next(f))
                except StopIteration:
                    break
                except UnicodeDecodeError:
                    head.append('UnicodeDecodeError\n')
            return f'Here are the first few lines of the file:\n' \
                   f'```output\n{"".join(head)}\n```\n'

    def pretty_repr(self, num_lines: int = 4, view_purpose: ViewPurpose = None, file_num: Optional[int] = None) -> str:
        if file_num is not None:
            s = f'### File {file_num}: "{self.file_path}"\n'
        else:
            s = f'### "{self.file_path}"\n'
        description = self.description
        if description is not None:
            description = hypertarget_if_referencable_text_product(description, view_purpose, with_header=False)
            s += f'{description}\n\n'
        if num_lines > 0 and not self.get_is_binary():
            s += self.get_file_header(num_lines)
        return s


class DataFileDescriptions(List[DataFileDescription]):
    """
    A list of data file descriptions.
    """

    def __init__(self, *args, data_folder: Optional[Union[str, Path]] = None,
                 general_description: Optional[Union[str, ReferencableTextProduct]] = None,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.data_folder = data_folder
        self.general_description = general_description

    @classmethod
    def from_other(cls, other: DataFileDescriptions):
        return cls(other, data_folder=other.data_folder, general_description=other.general_description)

    def __str__(self):
        return self.pretty_repr()

    def get_file_description(self, file_path: str) -> DataFileDescription:
        """
        Return the data file description for the given file path.
        """
        for data_file_description in self:
            if data_file_description.file_path == file_path:
                return data_file_description
        raise ValueError(f"Could not find file path {file_path} in data file descriptions.")

    def get_children(self, data_file: DataFileDescription):
        """
        Return the children of the given data file.
        """
        return DataFileDescriptions([data_file_description for data_file_description in self
                                     if data_file_description.originated_from == data_file.file_path],
                                    data_folder=self.data_folder)

    def get_all_raw_files(self):
        """
        Return all the raw files.
        """
        return DataFileDescriptions([data_file_description for data_file_description in self
                                     if data_file_description.originated_from is None],
                                    data_folder=self.data_folder)

    def get_pretty_description_for_file_and_children(
            self, data_file: DataFileDescription, index: Mutable = None, view_purpose: ViewPurpose = None):
        """
        Return a pretty description for the given data file and all its children.
        """
        children = self.get_children(data_file)
        index.val += 1
        s = data_file.pretty_repr(0 if children else 4, view_purpose=view_purpose, file_num=index.val) + '\n'
        for child in children:
            s += self.get_pretty_description_for_file_and_children(child, index)
        return s

    def pretty_repr(self, num_lines: int = 4, view_purpose: ViewPurpose = None) -> str:
        s = ''
        if self.general_description is not None:
            s += '## General Description\n'
            s += hypertarget_if_referencable_text_product(self.general_description, view_purpose,
                                                          with_header=False) + '\n'
        with run_in_directory(self.data_folder):
            s += '## Data Files\n'
            if len(self) == 0:
                s += 'There are no data files'
            elif len(self) == 1:
                s += f"The dataset consists of 1 data file:\n\n"
                s += self[0].pretty_repr(num_lines, view_purpose=view_purpose)
            else:
                s += f"The dataset consists of {len(self)} data files:\n\n"
                index = Mutable(0)
                for parent in self.get_all_raw_files():
                    s += self.get_pretty_description_for_file_and_children(parent, index, view_purpose=view_purpose)
            return s

    def get_data_filenames(self):
        return [data_file_description.file_path for data_file_description in self]

    def to_latex(self,
                 section_name: str = 'Data Description',
                 text: str = 'Here is the data description, as provided by the user:',
                 view_purpose: ViewPurpose = None) -> str:
        s = ''
        s += f"\\section{{{section_name}}}\n{text}\n\n"
        s += wrap_as_latex_code_output(
            self.pretty_repr(num_lines=0, view_purpose=view_purpose))
        return s


@dataclass
class CreateDataFileDescriptions:
    """
    Given a project directory, create a temp folder to run in, and
    the data file descriptions.
    """
    project_directory: Path = None
    data_files_str_paths: List[str] = field(default_factory=list)
    data_files_is_binary: List[Optional[bool]] = field(default_factory=list)

    GENERAL_DESCRIPTION_FILENAME = 'general_description.txt'
    DESCRIPTION_FILENAME_EXT = '.description.txt'

    GENERAL_FILE_DESCRIPTION_PREFIX = 'S'
    FILE_DESCRIPTIONS_PREFIXES = ('T', 'U', 'V', 'W', 'X', 'Y', 'Z')

    temp_folder_to_run_in: Path = FOLDER_FOR_RUN

    def _get_description_file_path(self, data_file_path_str: str):
        data_file_path = self._convert_data_file_path_str_to_path(data_file_path_str)
        return self.project_directory / (data_file_path.name + self.DESCRIPTION_FILENAME_EXT)

    def _read_file_description(self, data_file_path_str: str):
        description_file_path = self._get_description_file_path(data_file_path_str)
        if not description_file_path.exists():
            raise FileNotFoundError(f"Description file for {data_file_path_str} not found.")
        return description_file_path.read_text()

    def _read_general_description(self):
        general_description_file_path = self.project_directory / self.GENERAL_DESCRIPTION_FILENAME
        if not general_description_file_path.exists():
            raise FileNotFoundError(f"General description file not found:\n{general_description_file_path}.")
        return general_description_file_path.read_text()

    def _get_hypertarget_prefix(self, file_num: Optional[int] = None, file_name: str = None):
        """
        Return the hypertarget prefix for the given file number.
        file_num: the file number. None refers to the general description.
        """
        return self.GENERAL_FILE_DESCRIPTION_PREFIX if file_num is None \
            else self.FILE_DESCRIPTIONS_PREFIXES[file_num]

    def _convert_description_to_referenceable_text(self, description: str, file_num: Optional[int],
                                                   file_name: str = None):
        """
        file_num: the file number. None refers to the general description.
        """
        return ReferencableTextProduct(
            referencable_text=NumericReferenceableText(
                text=unidecode(description),
                hypertarget_prefix=self._get_hypertarget_prefix(file_num, file_name)),
            name=file_name,
        )

    def _convert_data_file_path_str_to_path(self, data_file_path: str):
        """
        convert to absolute path.
        if relative path, it is relative to the project directory.
        """
        return Path(data_file_path).absolute() if Path(data_file_path).is_absolute() \
            else self.project_directory / data_file_path

    def _get_file_description_referenceable_text(self, data_file_str_path: str, file_num: Optional[int]):
        description = self._read_file_description(data_file_str_path)
        data_file_path = self._convert_data_file_path_str_to_path(data_file_str_path)
        return self._convert_description_to_referenceable_text(description, file_num, data_file_path.name)

    def _get_general_description_referenceable_text(self):
        return self._convert_description_to_referenceable_text(self._read_general_description(), None,
                                                               'General Description')

    def _copy_files_and_get_list_of_data_file_descriptions(self) -> List[DataFileDescription]:
        data_file_descriptions = []
        clear_directory(self.temp_folder_to_run_in)  # clear data folder
        for j, data_file_str_path in enumerate(self.data_files_str_paths):
            data_file_path = self._convert_data_file_path_str_to_path(data_file_str_path)
            data_file_path_zip = data_file_path.with_name(data_file_path.name + '.zip')
            if os.path.exists(data_file_path):
                # copy file to data folder
                shutil.copyfile(data_file_path, self.temp_folder_to_run_in / data_file_path.name)
            elif os.path.exists(data_file_path_zip):
                # unzip file to data folder
                with zipfile.ZipFile(data_file_path_zip, 'r') as zip_ref:
                    zip_ref.extractall(self.temp_folder_to_run_in)
            else:
                raise FileNotFoundError(f"File {data_file_path.name} or {data_file_path.name}.zip "
                                        f"not found in {data_file_path.parent}")

            data_file_descriptions.append(DataFileDescription(
                file_path=data_file_path.name,  # relative to data dir
                description=self._get_file_description_referenceable_text(data_file_str_path, j),
                is_binary=self.data_files_is_binary[j]
            ))
        return data_file_descriptions

    def create_temp_folder_and_get_file_descriptions(self) -> DataFileDescriptions:
        file_descriptions = self._copy_files_and_get_list_of_data_file_descriptions()
        return DataFileDescriptions(
            file_descriptions,
            data_folder=self.temp_folder_to_run_in,
            general_description=self._get_general_description_referenceable_text()
        )

    def get_raw_str_data_file_descriptions(self) -> DataFileDescriptions:
        return DataFileDescriptions(
            [DataFileDescription(
                file_path=data_file_str_path,  # relative to data dir
                description=self._read_file_description(data_file_str_path),
                is_binary=self.data_files_is_binary[j]
            ) for j, data_file_str_path in enumerate(self.data_files_str_paths)],
            data_folder=None,
            general_description=self._read_general_description(),
        )

    def create_file_descriptions(self, general_description: str, data_file_descriptions: List[str]):
        """
        Create the file descriptions.
        """
        (self.project_directory / self.GENERAL_DESCRIPTION_FILENAME).write_text(general_description)
        for j, data_file_str_path in enumerate(self.data_files_str_paths):
            self._get_description_file_path(data_file_str_path).write_text(data_file_descriptions[j])

    def check_files_exist(self):
        for data_file_str_path in self.data_files_str_paths:
            data_file_path = self._convert_data_file_path_str_to_path(data_file_str_path)
            data_file_path_zip = data_file_path.with_name(data_file_path.name + '.zip')
            if not data_file_path.exists() and not data_file_path_zip.exists():
                raise FileNotFoundError(f"Data file {data_file_path} not found.")
