import dataclasses
import gzip
import json
import logging
import os
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Optional

import requests

from jf_ingest import logging_helper
from jf_ingest.config import IngestionConfig, IngestionType
from jf_ingest.utils import retry_session

logger = logging.getLogger(__name__)


class SubDirectory:
    JIRA = "jira"
    GIT = "git"


class StrDefaultEncoder(json.JSONEncoder):
    def default(self, o):
        if dataclasses.is_dataclass(o):
            return dataclasses.asdict(o)
        elif type(o) == datetime:
            return o.isoformat()
        return str(o)


class IngestIOHelper:
    def __init__(self, ingest_config: IngestionConfig):
        self.ingest_config = ingest_config
        # EVERYTHING in this file path will (potentially) be uploaded to S3!
        # DO NOT put any creds file in this path!!!!
        self.local_file_path = ingest_config.local_file_path

        if not os.path.exists(self.local_file_path):
            os.makedirs(self.local_file_path)

    def _get_file_name(self, object_name: str, batch_number: Optional[int] = 0):
        return f'{object_name}{batch_number if batch_number else ""}.json'

    def _serialize_json_to_bytes(
        self, json_data: dict | list[dict], indent_level: int = 2
    ) -> bytes:
        return json.dumps(json_data, indent=indent_level, cls=StrDefaultEncoder).encode("utf-8")

    def _write_file(self, json_data: dict | list[dict], full_file_path: str):
        """
        Writes json file data to local file system at filepath. Called by write_json_data_to_local.
        :param json_data:
        :param full_file_path:
        :return: nothing
        """
        with open(full_file_path, "wb") as f:
            f.write(self._serialize_json_to_bytes(json_data=json_data))
            logger.debug(f"File: {full_file_path}, Size: {round(f.tell() / 1000000, 1)}MB")

    def write_json_to_local_or_s3(
        self,
        object_name: str,
        json_data: dict | list[dict],
        subdirectory: SubDirectory,
        save_locally: bool,
        upload_to_s3: bool,
        git_instance_key: str = None,
        batch_number: Optional[int] = 0,
    ) -> None:
        """Helper function for writing data locally or pushing data to S3.
        NOTE: When writing git data, we require an instance key. This is because
        we can have multiple git instances, so we rely on a uniqueness key to keep
        things separated. The key is generated by the GithubInstance object in Jellyfish

        Args:
            object_name (str): The object string we are writing
            json_data (dict | list[dict]): The 'raw' data
            subdirectory (SubDirectory): The directory we are pushing to (JIRA or GIT)
            save_locally (bool): If true, save locally
            upload_to_s3 (bool): If true, push to S3 (remote)
            git_instance_key (str): When provided, adds in an additional suffix to the subdirectory. Required for Git!
            batch_number (Optional[int], optional): Used for breaking up very large files, like JiraIssues. Defaults to 0.
        """
        if subdirectory == SubDirectory.GIT:
            if git_instance_key:
                subdirectory_str = f'{subdirectory}_{git_instance_key}'
            else:
                raise Exception(
                    f'When writing Git Data with this function, the file_prefix argument is required!'
                )
        else:
            subdirectory_str = str(subdirectory)

        if save_locally:
            self.write_json_data_to_local(
                object_name=object_name,
                json_data=json_data,
                subdirectory=subdirectory_str,
                batch_number=batch_number,
            )
        if upload_to_s3:
            self.write_json_data_to_s3(
                object_name=object_name,
                json_data=json_data,
                subdirectory=subdirectory_str,
                batch_number=batch_number,
            )

    def write_json_data_to_local(
        self,
        object_name: str,
        json_data: dict | list[dict],
        subdirectory: str,
        batch_number: Optional[int] = 0,
    ) -> None:
        """
        Writes json data to local file system in proper subdirectory (depending on jira or git)
        :param object_name: from JiraObject or GitObject enum, eg "jira_boards", "git_prs", etc
        :param json_data: data returned from ingestion
        :param subdirectory: directory under `<output_dir>/<timestamp>/` to write to, eg "jira" or "git_EXAMPLEKEY"
        :param batch_number: if this is a batched file, the batch number
        :return: None
        """
        file_name = self._get_file_name(object_name=object_name, batch_number=batch_number)
        full_file_path = f"{self.local_file_path}/{subdirectory}/{file_name}"
        logger.debug(f"Attempting to save {object_name} data to {full_file_path}")
        try:
            self._write_file(json_data, full_file_path)
        except FileNotFoundError:
            Path(full_file_path).parent.mkdir(exist_ok=True, parents=True)
            self._write_file(json_data, full_file_path)
        logger.debug(f"Successfully saved {object_name} data to {full_file_path}")

    def write_json_data_to_s3(
        self,
        object_name: str,
        json_data: dict | list[dict],
        subdirectory: str,
        batch_number: Optional[int] = 0,
    ) -> None:
        """Helper function for uploading JSON Data to S3, one file at a time.
        Called at potentially the same time that we save data to the local disk.
        We want to continuously upload files to S3 and potentially not save them
        locally to reduce our disk footprint when working in Dask Tasks, which
        are limited to 20 GB as of January 2024

        Args:
            object_name (str): The object string we are writing
            json_data (dict | list[dict]): The "raw" data
            subdirectory (str): The subdirectory we are pushing to (jira/ or git_EXAMPLEKEY/)
            batch_number (Optional[int], optional): Used for breaking up very large files, like JiraIssues. Defaults to 0.
        """
        # Get signed URL for this file
        filename = f"{subdirectory}/{self._get_file_name(object_name=object_name, batch_number=batch_number)}"
        s3_path, url, fields = self.get_signed_url(filename=filename)

        encoded_json_data = self._serialize_json_to_bytes(json_data=json_data, indent_level=0)
        logging_helper.send_to_agent_log_file(f"Compressing {object_name} data...")
        zipped_json_data = gzip.compress(encoded_json_data)
        session = retry_session()
        # When pushing to S3, there's a small chance that we randomly get a 500 error.
        # It is expected that on retry this should be resolved.
        # To be safe, we've implemented an exponential back off algorithm
        max_attempts = 5
        current_attempt = 0
        while True:
            try:
                logging_helper.send_to_agent_log_file(
                    f"Attempting to upload {object_name} data to S3"
                )
                upload_resp = session.post(
                    url,
                    data=fields,
                    files={"file": (s3_path, zipped_json_data)},
                )
                upload_resp.raise_for_status()
                logging_helper.send_to_agent_log_file(
                    f"Successfully uploaded {object_name} data to S3"
                )
                break
            except Exception as e:
                # Increment attempt number
                current_attempt += 1
                logging_helper.send_to_agent_log_file(
                    f"Exception encountered when attempting to upload {object_name} data to S3. Attempt number {current_attempt} of {max_attempts} Error: {e}",
                    level=logging.ERROR,
                )
                if current_attempt >= max_attempts:
                    logger.error(f'Error attempting to upload {object_name} data to S3')
                    logging_helper.send_to_agent_log_file(
                        traceback.format_exc(), level=logging.ERROR
                    )
                    raise
                else:
                    sleep_time = current_attempt**2
                    logging_helper.send_to_agent_log_file(f'Sleeping for {sleep_time}')
                    time.sleep(sleep_time)

        return

    def get_signed_url(self, filename: str) -> None:
        """
        Gets signed url from Jellyfish API for uploading to S3.
        This hits an API that is designed to return multiple signed URLs,
        but we are using it for just one upload
        :param filename: list of filenames that will need urls
        :param timestamp: timestamp created at beginning of ingestion (used for output directory)
        :return: A dictionary in the following shape:
        {
            'signed_urls: {
                <filename>: {
                    's3_path': '<s3_path>',
                    'url': {
                        'url': '<signed_url>',
                        'fields': {
                            '<field>': '<value>' for field in (key, AWSAccessKeyId, x-amz-security-token, policy, signature)
                        }
                    }
                },
            },
        }
        """
        filename = f"{filename}.gz"
        base_url = self.ingest_config.jellyfish_api_base
        headers = {"Jellyfish-API-Token": self.ingest_config.jellyfish_api_token}
        payload = {
            # NOTE: The API handles a LIST of file names, so we need to wrap this
            # individual file in brackets
            "files": [filename],
            "ingestType": self.ingest_config.ingest_type,
        }  # the files will be gzipped before uploading

        r = requests.post(
            f"{base_url}/endpoints/ingest/signed-url?timestamp={self.ingest_config.timestamp}",
            headers=headers,
            json=payload,
        )
        r.raise_for_status()
        json_resp = r.json()["signed_urls"][filename]

        return json_resp['s3_path'], json_resp['url']['url'], json_resp['url']['fields']

    def get_signed_urls(
        self, filenames: list[str], timestamp: str, ingestion_type: IngestionType
    ) -> dict[str, dict]:
        """
        Gets signed urls from Jellyfish API for uploading to S3
        :param filenames: list of filenames that will need urls
        :param timestamp: timestamp created at beginning of ingestion (used for output directory)
        :return: the signed urls dict of shape:
        {
            'signed_urls: {
                <filename>: {
                    's3_path': '<s3_path>',
                    'url': {
                        'url': '<signed_url>',
                        'fields': {
                            '<field>': '<value>' for field in (key, AWSAccessKeyId, x-amz-security-token, policy, signature)
                        }
                    }
                },
            },
        }
        """
        base_url = self.ingest_config.jellyfish_api_base
        headers = {"Jellyfish-API-Token": self.ingest_config.jellyfish_api_token}
        payload = {
            "files": [f"{name}.gz" for name in filenames],
            "ingestType": ingestion_type,
        }  # the files will be gzipped before uploading

        r = requests.post(
            f"{base_url}/endpoints/ingest/signed-url?timestamp={timestamp}",
            headers=headers,
            json=payload,
        )
        r.raise_for_status()

        return r.json()["signed_urls"]
