import os
import shutil
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, List, Optional

from huggingface_hub import Repository
from loguru import logger
from prettytable import PrettyTable

from .splits import TEST_SPLIT, TRAIN_SPLIT, VALID_SPLIT
from .tasks import TASKS
from .utils import BOLD_TAG, CYAN_TAG, GREEN_TAG, PURPLE_TAG, RESET_TAG, YELLOW_TAG, http_get, http_post
from .validation import validate_file


FILE_STATUS = (
    "☁ Uploaded",
    "⌚ Queued",
    "⚙ In Progress...",
    "✅ Success!",
    "❌ Failed: file not found",
    "❌ Failed: unsupported file type",
    "❌ Failed: server error",
    "❌ Invalid column mapping, please fix it and re-upload the file.",
)

JOB_STATUS = (
    ("⌚", "queued"),
    ("🚀", "start"),
    ("⚙", "data_munging"),
    ("🏃", "model_training"),
    ("✅", "success"),
    ("❌", "failed"),
)

PROJECT_STATUS = (
    ("✨", "Created"),
    ("🚀", "Data processing started"),
    ("✅", "Data processing successful"),
    ("❌", "Failed to download data files from the huggingface hub"),
    ("❌", "Missing 'train' or 'valid' split in data files"),
    ("❌", "Failed to process data files"),
    ("❌", "Failed to upload processed data files to the huggingface hub"),
)


SPLITS = (TRAIN_SPLIT, VALID_SPLIT, TEST_SPLIT)


@dataclass
class TrainingJob:
    """A training job in AutoNLP"""

    job_id: int
    status: str
    status_emoji: str
    created_at: datetime
    updated_at: datetime

    @classmethod
    def from_json_resp(cls, json_resp: dict):
        return cls(
            job_id=json_resp["id"],
            status_emoji=JOB_STATUS[json_resp["status"] - 1][0],
            status=JOB_STATUS[json_resp["status"] - 1][1],
            created_at=datetime.fromisoformat(json_resp["created_at"]),
            updated_at=datetime.fromisoformat(json_resp["updated_at"]),
        )

    def __str__(self):
        return "\n".join(
            [
                f"📚 Model # {self.job_id}",
                f"   • {BOLD_TAG}Status{RESET_TAG}:      {self.status_emoji} {self.status}",
                f"   • {BOLD_TAG}Created at{RESET_TAG}:  {self.created_at.strftime('%Y-%m-%d %H:%M Z')}",
                f"   • {BOLD_TAG}Last update{RESET_TAG}: {self.updated_at.strftime('%Y-%m-%d %H:%M Z')}",
            ]
        )


@dataclass
class UploadedFile:
    """A file uploaded to an AutoNLP project"""

    file_id: int
    filename: str
    processing_status: str
    split: str
    col_mapping: Dict[str, str]
    created_at: datetime
    updated_at: datetime

    @classmethod
    def from_json_resp(cls, json_resp: dict):
        return cls(
            file_id=json_resp["data_file_id"],
            filename=json_resp["fname"],
            processing_status=FILE_STATUS[json_resp["download_status"] - 1],
            split=SPLITS[json_resp["split"] - 1],
            col_mapping=json_resp["col_mapping"],
            created_at=datetime.fromisoformat(json_resp["created_at"]),
            updated_at=datetime.fromisoformat(json_resp["updated_at"]),
        )

    def __str__(self):
        return "\n".join(
            [
                f"📁 {CYAN_TAG}{self.filename}{RESET_TAG} (id # {self.file_id})",
                f"   • {BOLD_TAG}Split{RESET_TAG}:             {self.split}",
                f"   • {BOLD_TAG}Processing status{RESET_TAG}: {self.processing_status}",
                f"   • {BOLD_TAG}Last update{RESET_TAG}:       {self.updated_at.strftime('%Y-%m-%d %H:%M Z')}",
            ]
        )


@dataclass
class Project:
    """An AutoNLP project"""

    _token: str
    proj_id: int
    name: str
    user: str
    task: str
    status_emoji: str
    status: str
    language: str
    created_at: datetime
    updated_at: datetime
    dataset_id: str
    files: Optional[List[UploadedFile]] = None
    training_jobs: Optional[List] = None

    @classmethod
    def from_json_resp(cls, json_resp: dict, token: str):
        """Build a Project from the API response, JSON-encoded"""
        return cls(
            proj_id=json_resp["id"],
            name=json_resp["proj_name"],
            user=json_resp["username"],
            task=list(filter(lambda key: TASKS[key] == json_resp["task"], TASKS.keys()))[0],
            status_emoji=PROJECT_STATUS[json_resp["status"] - 1][0],
            status=PROJECT_STATUS[json_resp["status"] - 1][1],
            created_at=datetime.fromisoformat(json_resp["created_at"]),
            updated_at=datetime.fromisoformat(json_resp["updated_at"]),
            dataset_id=json_resp["dataset_id"],
            language=json_resp["config"]["language"],
            _token=token,
        )

    def refresh(self):
        """Update information about uploaded files and models attached to the project"""
        logger.info("🔄 Refreshing uploaded files information...")
        resp = http_get(path=f"/projects/{self.proj_id}/data", token=self._token)
        json_files = resp.json()
        self.files = [UploadedFile.from_json_resp(file) for file in json_files]

        logger.info("🔄 Refreshing models information...")
        resp = http_get(path=f"/projects/{self.proj_id}/jobs", token=self._token)
        json_jobs = resp.json()
        self.training_jobs = [TrainingJob.from_json_resp(job) for job in json_jobs]

    def upload(self, filepaths: List[str], split: str, col_mapping: Dict[str, str]):
        """Uploads files to the project"""
        local_dataset_dir = os.path.expanduser(f"~/.huggingface/autonlp/projects/{self.dataset_id}")
        if os.path.exists(local_dataset_dir):
            if os.path.isdir(os.path.join(local_dataset_dir, "git")):
                clone_from = None
            else:
                shutil.rmtree(local_dataset_dir)
                clone_from = "https://huggingface.co/datasets/" + self.dataset_id
        else:
            clone_from = "https://huggingface.co/datasets/" + self.dataset_id
        dataset_repo = Repository(
            local_dir=local_dataset_dir,
            clone_from=clone_from,
            use_auth_token=self._token,
        )
        dataset_repo.git_pull()

        for idx, file_path in enumerate(filepaths):
            if not os.path.isfile(file_path):
                logger.error(f"[{idx + 1}/{len(filepaths)}] ❌ '{file_path}' does not exist or is not a file!")
                continue
            file_name = os.path.basename(file_path)
            file_extension = file_name.split(".")[-1]
            src = os.path.expanduser(file_path)
            dst = os.path.join(local_dataset_dir, "raw", file_name)
            logger.info(f"[{idx + 1}/{len(filepaths)}] 📦 Copying {src} to {dst}...")
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            shutil.copyfile(src, dst)

            logger.info(f"[{idx + 1}/{len(filepaths)}] 🔎 Validating {dst} and column mapping...")
            validate_file(path=dst, task=self.task, file_ext=file_extension, col_mapping=col_mapping)

            dataset_repo.lfs_track(patterns=[f"raw/*.{file_extension}"])

        dataset_repo.git_pull()

        try:
            logger.info("☁ Uploading files to the dataset hub...")
            dataset_repo.push_to_hub(commit_message="Upload from AutoNLP CLI")
            logger.info("✅ Successfully uploaded  the files!")
        except OSError as err:
            if "nothing to commit, working tree clean" in err.args[0]:
                logger.info("❔ Files did not change since last upload!")
                dataset_repo.git_push()
                return
            logger.error("❌ Something went wrong when uploading the files!")
            raise

        for idx, file_path in enumerate(filepaths):
            file_name = os.path.basename(file_path)
            logger.info(f"[{idx + 1}/{len(filepaths)}] 📁 Registering file {file_name} into project '{file_name}'...")
            payload = {
                "split": split,
                "col_mapping": col_mapping,
                "data_files": [{"fname": file_name, "username": self.user}],
            }
            http_post(path=f"/projects/{self.proj_id}/data/add", payload=payload, token=self._token)
            logger.info(f"[{idx + 1}/{len(filepaths)}] ✅ Success!")

    def train(self):
        """Starts training on the models"""
        http_get(path=f"/projects/{self.proj_id}/data/start_process", token=self._token)
        logger.info("🔥🔥 Training started!")

    def __str__(self):
        header = "\n".join(
            [
                f"AutoNLP Project (id # {self.proj_id})",
                "~" * 35,
                f" • {BOLD_TAG}Name{RESET_TAG}:        {PURPLE_TAG}{self.name}{RESET_TAG}",
                f" • {BOLD_TAG}Owner{RESET_TAG}:       {GREEN_TAG}{self.user}{RESET_TAG}",
                f" • {BOLD_TAG}Status{RESET_TAG}:      {BOLD_TAG}{self.status_emoji} {self.status}{RESET_TAG}",
                f" • {BOLD_TAG}Task{RESET_TAG}:        {YELLOW_TAG}{self.task.title().replace('_', ' ')}{RESET_TAG}",
                f" • {BOLD_TAG}Created at{RESET_TAG}:  {self.created_at.strftime('%Y-%m-%d %H:%M Z')}",
                f" • {BOLD_TAG}Last update{RESET_TAG}: {self.updated_at.strftime('%Y-%m-%d %H:%M Z')}",
                "",
            ]
        )
        printout = [header]

        # Uploaded files information
        if self.files is None:
            descriptions = ["❓ Files information unknown, update the project"]
        else:
            if len(self.files) == 0:
                descriptions = ["🤷 No files uploaded yet!"]
            else:
                sorted_files = sorted(self.files, key=lambda file: file.split)  # Sort by split
                descriptions = [str(file) for file in sorted_files]
        printout.append(
            "\n".join(
                [
                    "~" * 14 + f" {BOLD_TAG}Files{RESET_TAG} " + "~" * 14,
                    "",
                    "Dataset ID:",
                    f"{CYAN_TAG}{self.dataset_id}{RESET_TAG}",
                    "",
                ]
                + descriptions
            )
        )

        # Training jobs information
        if self.training_jobs is None:
            jobs_str = "❓ Models information unknown, update the project"
        else:
            if len(self.training_jobs) == 0:
                jobs_str = "🤷 No train jobs started yet!"
            else:
                model_table = PrettyTable(["", "ID", "Status", "Creation date", "Last update"])
                for job in sorted(self.training_jobs, key=lambda job: job.job_id):
                    model_table.add_row(
                        [
                            job.status_emoji,
                            job.job_id,
                            job.status,
                            job.created_at.strftime("%Y-%m-%d %H:%M Z"),
                            job.updated_at.strftime("%Y-%m-%d %H:%M Z"),
                        ]
                    )
                jobs_str = str(model_table)
        printout.append("\n".join(["", "~" * 12 + f" {BOLD_TAG}Models{RESET_TAG} " + "~" * 11, "", jobs_str]))

        return "\n".join(printout)

[((154, 8, 154, 68), 'loguru.logger.info', 'logger.info', ({(154, 20, 154, 67): '"""🔄 Refreshing uploaded files information..."""'}, {}), "('🔄 Refreshing uploaded files information...')", False, 'from loguru import logger\n'), ((159, 8, 159, 60), 'loguru.logger.info', 'logger.info', ({(159, 20, 159, 59): '"""🔄 Refreshing models information..."""'}, {}), "('🔄 Refreshing models information...')", False, 'from loguru import logger\n'), ((166, 28, 166, 100), 'os.path.expanduser', 'os.path.expanduser', ({(166, 47, 166, 99): 'f"""~/.huggingface/autonlp/projects/{self.dataset_id}"""'}, {}), "(f'~/.huggingface/autonlp/projects/{self.dataset_id}')", False, 'import os\n'), ((167, 11, 167, 44), 'os.path.exists', 'os.path.exists', ({(167, 26, 167, 43): 'local_dataset_dir'}, {}), '(local_dataset_dir)', False, 'import os\n'), ((175, 23, 179, 9), 'huggingface_hub.Repository', 'Repository', (), '', False, 'from huggingface_hub import Repository\n'), ((227, 8, 227, 49), 'loguru.logger.info', 'logger.info', ({(227, 20, 227, 48): '"""🔥🔥 Training started!"""'}, {}), "('🔥🔥 Training started!')", False, 'from loguru import logger\n'), ((186, 24, 186, 51), 'os.path.basename', 'os.path.basename', ({(186, 41, 186, 50): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((188, 18, 188, 47), 'os.path.expanduser', 'os.path.expanduser', ({(188, 37, 188, 46): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((189, 18, 189, 67), 'os.path.join', 'os.path.join', ({(189, 31, 189, 48): 'local_dataset_dir', (189, 50, 189, 55): '"""raw"""', (189, 57, 189, 66): 'file_name'}, {}), "(local_dataset_dir, 'raw', file_name)", False, 'import os\n'), ((192, 12, 192, 37), 'shutil.copyfile', 'shutil.copyfile', ({(192, 28, 192, 31): 'src', (192, 33, 192, 36): 'dst'}, {}), '(src, dst)', False, 'import shutil\n'), ((202, 12, 202, 68), 'loguru.logger.info', 'logger.info', ({(202, 24, 202, 67): '"""☁ Uploading files to the dataset hub..."""'}, {}), "('☁ Uploading files to the dataset hub...')", False, 'from loguru import logger\n'), ((204, 12, 204, 64), 'loguru.logger.info', 'logger.info', ({(204, 24, 204, 63): '"""✅ Successfully uploaded  the files!"""'}, {}), "('✅ Successfully uploaded  the files!')", False, 'from loguru import logger\n'), ((214, 24, 214, 51), 'os.path.basename', 'os.path.basename', ({(214, 41, 214, 50): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((67, 23, 67, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(67, 46, 67, 69): "json_resp['created_at']"}, {}), "(json_resp['created_at'])", False, 'from datetime import datetime\n'), ((68, 23, 68, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(68, 46, 68, 69): "json_resp['updated_at']"}, {}), "(json_resp['updated_at'])", False, 'from datetime import datetime\n'), ((102, 23, 102, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(102, 46, 102, 69): "json_resp['created_at']"}, {}), "(json_resp['created_at'])", False, 'from datetime import datetime\n'), ((103, 23, 103, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(103, 46, 103, 69): "json_resp['updated_at']"}, {}), "(json_resp['updated_at'])", False, 'from datetime import datetime\n'), ((145, 23, 145, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(145, 46, 145, 69): "json_resp['created_at']"}, {}), "(json_resp['created_at'])", False, 'from datetime import datetime\n'), ((146, 23, 146, 70), 'datetime.datetime.fromisoformat', 'datetime.fromisoformat', ({(146, 46, 146, 69): "json_resp['updated_at']"}, {}), "(json_resp['updated_at'])", False, 'from datetime import datetime\n'), ((168, 29, 168, 67), 'os.path.join', 'os.path.join', ({(168, 42, 168, 59): 'local_dataset_dir', (168, 61, 168, 66): '"""git"""'}, {}), "(local_dataset_dir, 'git')", False, 'import os\n'), ((171, 16, 171, 48), 'shutil.rmtree', 'shutil.rmtree', ({(171, 30, 171, 47): 'local_dataset_dir'}, {}), '(local_dataset_dir)', False, 'import shutil\n'), ((183, 19, 183, 44), 'os.path.isfile', 'os.path.isfile', ({(183, 34, 183, 43): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((191, 24, 191, 44), 'os.path.dirname', 'os.path.dirname', ({(191, 40, 191, 43): 'dst'}, {}), '(dst)', False, 'import os\n'), ((210, 12, 210, 78), 'loguru.logger.error', 'logger.error', ({(210, 25, 210, 77): '"""❌ Something went wrong when uploading the files!"""'}, {}), "('❌ Something went wrong when uploading the files!')", False, 'from loguru import logger\n'), ((274, 30, 274, 95), 'prettytable.PrettyTable', 'PrettyTable', ({(274, 42, 274, 94): "['', 'ID', 'Status', 'Creation date', 'Last update']"}, {}), "(['', 'ID', 'Status', 'Creation date', 'Last update'])", False, 'from prettytable import PrettyTable\n'), ((207, 16, 207, 74), 'loguru.logger.info', 'logger.info', ({(207, 28, 207, 73): '"""❔ Files did not change since last upload!"""'}, {}), "('❔ Files did not change since last upload!')", False, 'from loguru import logger\n')]