import base64
import json
import os
import pathlib
import tempfile
from http import HTTPStatus

import backoff
import boto3
import requests
import rollbar
import xmltodict

from .exceptions import (
    ConstellationException,
    ConstellationInvalidMethodException,
    ConstellationInvalidResponseException,
    ConstellationMissingFieldException,
    ConstellationNotLoggedInException,
    ConstellationObjectDoesNotExistException,
)
from .models import CloudFileArgs, ConstellationFieldParser
from .utils import md5_hash_file, parts_to_xml


class Constellation:
    _UPLOAD_PART_SIZE = 10_000_000  # 10 MB
    _MULTIPART_CLOUDFILE_THRESHOLD = 100_000_000  # 100 MB

    def __init__(self, auth_token="", host=""):
        self._host = host or "https://constellation.emeraldcloudlab.com"
        self._auth_token = auth_token

        self._temp_directory = None

    # Retry decorator for requests failures (like timeout or connection error)
    _retry_requests_decorator = backoff.on_exception(
        wait_gen=backoff.expo,
        exception=(requests.exceptions.Timeout, requests.exceptions.ConnectionError),
        max_tries=5,
    )

    # Configuration related functions
    def host(self):
        """Return the url of the instance constellation that the client is connecting to"""
        return self._host

    def download_directory(self):
        """Returns where cloud file downloads are stored"""
        if self._temp_directory is None:
            self._temp_directory = tempfile.TemporaryDirectory()

        return self._temp_directory.name

    # Identity requests
    @_retry_requests_decorator
    def login(self, username, password):
        """Return the auth token token for the supplied username and password"""
        auth_token = self._send_and_validate_request(
            "/ise/signintoken",
            {"username": username, "password": password},
            "AuthToken",
            requires_auth_token=False,
        )
        self.set_auth_token(auth_token)
        return self.me()

    def is_logged_in(self):
        return self._auth_token != ""

    def set_auth_token(self, auth_token):
        """Set the auth token used by the constellation client"""
        self._auth_token = auth_token

    @_retry_requests_decorator
    def me(self):
        """Return the response from the /ise/me endpoint"""
        return self._send_and_validate_request("/ise/me", None, None, method="GET")

    # Searching objects
    def search(self, ecl_type, query):
        """Perform a simple search on a given type"""
        payload = {
            "queries": [
                {
                    "clauses": [
                        {
                            "Types": [ecl_type],
                            "Query": query,
                        }
                    ],
                    "SubTypes": True,
                    "date": None,
                    "ignore_time": False,
                }
            ],
        }

        return self._send_and_validate_request(
            "/obj/search",
            payload,
            field_to_retrieve="Results",
        )

    # Downloading objects
    def download(self, objects, field_string="All", auto_download_cloud_files=False):
        """Download the supplied fields off the supplied object or objects.

        Note that fields can traverse across links using []'s.  For example:

        download(self.me()["ID"], "FinancingTeams[Name]")

        will download the status field off the protocol link of the supplied
        object.
        """
        result = None
        collapse_fields = False
        collapse_objects = False

        # Normalize the object and fields so they're always lists,
        # even if they're passed as single things.  This allows to
        # handle things like download("obj1", "Name") and
        # download([obj1], ["Name"]) in the same interface
        if not isinstance(field_string, list):
            field_string = [field_string]
            collapse_fields = True

        if not isinstance(objects, list):
            objects = [objects]
            collapse_objects = True

        result = self._download_objects(objects, field_string, auto_download_cloud_files)

        # Unwind the normalization of fields and object ids into lists
        if collapse_fields and collapse_objects:
            return result[0][0]
        if collapse_fields:
            return [x[0] for x in result]
        if collapse_objects:
            return result[0]

        return result

    def _download_objects(self, objects, field_string, auto_download_cloud_files):
        """Download multiple objects in a single request"""

        # Build the request
        traversals = self._parse_download_field_string(field_string)
        obj_requests = []
        for obj in objects:
            obj_request = {"object": {"id": obj.id}, "limit": 0, "subfield_info": True}
            if traversals:
                obj_request["traversals"] = traversals
            obj_requests.append(obj_request)

        # Send the request to constellation
        responses = self._send_and_validate_request(
            "/obj/download",
            {"requests": obj_requests},
            "responses",
        )
        if len(responses) != len(objects):
            raise ConstellationInvalidResponseException(
                f"Expected exactly {len(objects)} response from /obj/download", responses
            )

        # Parse the json results into objects
        response_objects = {}
        field_parser = ConstellationFieldParser(self.download_blob_ref)
        for response in responses:
            (
                sub_response_objects,
                sub_response_field_summaries,
            ) = self._parse_field_values_from_response(response)
            response_objects.update(sub_response_objects)
            for ecl_type, field_summaries in sub_response_field_summaries.items():
                field_parser.add_summaries_for_type(field_summaries, ecl_type)

        # Map the objects into a formatted result based on the fields string
        parsed_results = self._parse_results_from_response(response_objects, objects, traversals, field_parser)

        # If requested, go ahead and download all the cloud files encountered
        if auto_download_cloud_files:
            for cloud_file in field_parser.cloud_files_encountered():
                self.download_cloud_file(cloud_file)

        # return the results
        return parsed_results

    def _parse_field_values_from_response(self, response):
        """Parse out the field values contained in the response"""
        response_objects = {}
        field_summaries = {}
        response_fields = response.get("fields", {})
        response_field_summaries = response.get("summaries", {})
        ecl_type = response.get("resolved_object", {}).get("type")
        object_id = response.get("resolved_object", {}).get("id")

        if ecl_type is None or object_id is None:
            raise ConstellationObjectDoesNotExistException(response.get("object", {}).get("id"))

        response_objects[object_id] = {"type": ecl_type, "id": object_id}
        response_objects[object_id].update(response_fields)
        if ecl_type not in field_summaries:
            field_summaries[ecl_type] = {}
        field_summaries[ecl_type].update(response_field_summaries)
        for traversed_object in response.get("traversed_objects", []):
            (
                traversed_response_objects,
                traversed_field_summaries,
            ) = self._parse_field_values_from_response(traversed_object)
            response_objects.update(traversed_response_objects)
            field_summaries.update(traversed_field_summaries)

        return response_objects, field_summaries

    def _parse_download_field_string(self, field_string):
        """Parse the download field string into fields and traversals for constellation"""
        return [self._parse_traversal_from_field_string(traversal_string) for traversal_string in field_string]

    def _parse_traversal_from_field_string(self, field_string):
        """Parse a download traversal from the string format.

        E.g., turn FinancingTeams[Name] -> constellation consumable traversal
        """
        tokens = field_string.split("[", 1)
        if len(tokens) == 0:
            return None

        current_traversal = {"name": tokens[0].replace("]", "")}
        if len(tokens) == 1:
            return current_traversal

        current_traversal.update({"next": [self._parse_traversal_from_field_string(tokens[1])]})
        return current_traversal

    def _parse_results_from_response(self, response_objects, objects, traversals, field_parser):
        """Parse out the response for the requested fields and traversals"""
        results = []
        for obj in objects:
            results.append(self._parse_results_from_traversals(response_objects, obj, traversals, field_parser))
        return results

    def _parse_results_from_traversals(self, objects_info, obj, traversals, field_parser):
        """Parse out the response for the requested traversals"""
        object_info = objects_info.get(obj.id, {})
        results = []
        for traversal in traversals:
            if "name" not in traversal:
                continue
            field_name = traversal["name"]
            field_value = object_info.get(field_name)
            if field_name == "All":
                all_dict = {}
                for field_name, field_value in object_info.items():
                    all_dict[field_name] = field_parser.parse_field_value(
                        object_info.get("type"), field_name, field_value
                    )
                return [all_dict]

            if "next" in traversal:
                sub_traversals = traversal["next"]
                if isinstance(field_value, list):
                    results.append(
                        [
                            self._parse_results_from_traversals(
                                objects_info,
                                field_parser.to_link(field_obj),
                                sub_traversals,
                                field_parser,
                            )
                            for field_obj in field_value
                        ][0]
                    )
                else:
                    results.append(
                        self._parse_results_from_traversals(
                            objects_info,
                            field_parser.to_link(field_value),
                            sub_traversals,
                            field_parser,
                        )[0]
                    )
            else:
                results.append(field_parser.parse_field_value(object_info.get("type"), field_name, field_value))

        return results

    # Downloading cloud files
    def download_cloud_file(self, cloud_file):
        """Download the supplied cloud file to client.download_directory()"""
        file_name, extension, file_info = self.download(cloud_file, ["FileName", "FileType", "CloudFile"])
        extension = extension.replace('"', "")
        if not file_name:
            file_name = "EmeraldCloudFile_{}".format(cloud_file.id.replace(":", "_"))
        target_path = self._get_temp_file_path(file_name, extension)
        target_path = self._download_cloud_file(
            file_info["CloudFileId"], file_info["Bucket"], file_info["Key"], target_path
        )

        cloud_file.local_path = target_path
        return target_path

    def _get_temp_file_path(self, filename, extension):
        """Get a temp file path as close to filename as possible"""
        index = 0
        target_path = os.path.join(self.download_directory(), f"{filename}.{extension}")
        while os.path.exists(target_path):
            target_path = os.path.join(self.download_directory(), f"{filename}_{index}.{extension}")
            index += 1

        return target_path

    def blobsign_download(self, cloud_file_id, cloud_file_bucket, cloud_file_key):
        """Return the url to download the supplied cloud file"""
        return self._send_and_validate_request(
            "/blobsign/download",
            {
                "CloudFileId": cloud_file_id,
                "Bucket": cloud_file_bucket,
                "Key": cloud_file_key,
            },
            "Url",
        )

    def _download_cloud_file(self, cloud_file_id, cloud_file_bucket, cloud_file_key, target_path):
        """Download a cloud file based off the supplied identifiers"""
        url = self.blobsign_download(cloud_file_id, cloud_file_bucket, cloud_file_key)
        if not url:
            raise ConstellationMissingFieldException("When downloading cloud file", "url")

        # Now download the file from s3
        target_file = open(target_path, "wb")
        download_response = requests.get(url)
        target_file.write(download_response.content)
        target_file.close()

        return target_path

    def download_blob_ref(self, blob_ref, target_path=""):
        """Download a blob ref to the supplied target path"""
        target_path = target_path or self._get_temp_file_path(blob_ref.file_hash, "tmp")

        cloud_file_bucket = blob_ref.bucket
        cloud_file_key = blob_ref.key()
        if not cloud_file_bucket:
            raise ConstellationMissingFieldException("When downloading cloud file", "bucket")
        if not cloud_file_key:
            raise ConstellationMissingFieldException("When downloading cloud file", "hash")

        target_path = self._download_cloud_file("", cloud_file_bucket, cloud_file_key, target_path)
        blob_ref.local_path = target_path
        return target_path

    def upload(self, object_type, object_id, new_field_values):
        """Performs a very simple upload of a list of fields on a single object.

        Pass None for object_id to create a new object."""
        request_body = {
            "object": {"id": object_id, "type": object_type},
            "fields": new_field_values,
        }
        if object_id is None:
            # If you have no object id because you're creating a new object,
            # you need to format it slightly differently
            request_body = {"type": object_type, "fields": new_field_values}
        responses = self._send_and_validate_request(
            "/obj/upload", {"requests": [request_body]}, "responses", method="PUT"
        )

        if len(responses) != 1:
            raise ConstellationInvalidResponseException("Expected exactly 1 response from /obj/upload", responses)

        return responses[0]

    # Uploading cloud files
    def upload_cloud_file(self, file_path):
        """Uploads the given file to constellation"""
        md5 = md5_hash_file(file_path)
        extension = pathlib.Path(file_path).suffix
        if extension:
            extension = extension[1:]  # Remove the dot
        with open(file_path, "rb") as file:
            cloud_file_args = CloudFileArgs(
                file=file,
                name=pathlib.Path(file_path).stem,
                extension=extension,
                path=file_path,
                size=os.path.getsize(file_path),
                key=f"{md5.hexdigest()}.{extension}",
                content_md5=base64.b64encode(md5.digest()).decode(),
                content_disposition=f'attachment; filename="{file_path}"',
            )

            if cloud_file_args.size < self._MULTIPART_CLOUDFILE_THRESHOLD:  # 100MB
                return self._single_sign_and_upload_file_to_s3(cloud_file_args)
            return self._multi_part_sign_and_upload_file_to_s3(cloud_file_args)

    @_retry_requests_decorator
    def _single_sign_and_upload_file_to_s3(self, cloud_file_args):
        blobsign_response = self._blobsign("upload", cloud_file_args)
        headers = {
            "Content-Md5": cloud_file_args.content_md5,
            "Content-Disposition": cloud_file_args.content_disposition,
        }
        cloud_file_args.file.seek(0)
        s3_response = requests.put(
            url=blobsign_response["Url"],
            data=cloud_file_args.file,
            headers=headers,
        )
        if s3_response.status_code != HTTPStatus.OK:
            raise ConstellationException(
                f"Unexpected non-200 response from constellation endpoint {s3_response}: {s3_response.text}",
                s3_response.status_code,
            )
        return self._upload_cloud_file_to_constellation(blobsign_response["Key"], cloud_file_args)

    @_retry_requests_decorator
    def _upload_cloud_file_to_constellation(self, s3_key, cloud_file_args):
        cloud_field_values = {
            "$Type": "__JsonEmeraldCloudFile__",
            "Bucket": "emeraldsci-ecl-blobstore-stage",
            "CloudFileId": "None",
            "Key": s3_key,
        }
        fields = {
            "FileName": cloud_file_args.name,
            "FileSize": f'Quantity[{cloud_file_args.size}, "bytes"]',
            "FileType": f'"{cloud_file_args.extension}"',
            "CloudFile": cloud_field_values,
        }
        return self.upload("Object.EmeraldCloudFile", None, fields)

    def _blobsign(self, sign_type, cloud_file_args: CloudFileArgs, part=0, upload_id=""):
        """Return the url to upload the supplied cloud file"""
        sign_response = self._send_and_validate_request(
            "/blobsign/" + sign_type,
            {
                "Key": cloud_file_args.key,
                "ContentMD5": cloud_file_args.content_md5,
                "ContentDisposition": cloud_file_args.content_disposition,
                "Part": part,
                "UploadID": upload_id,
            },
            additional_headers={"ContentDisposition": cloud_file_args.content_disposition},
        )
        return sign_response

    def _multi_part_sign_and_upload_file_to_s3(self, cloud_file_args):
        self.s3_client = boto3.client("s3")
        upload_id, sharded_key = self._send_create_multi_part_upload_request(cloud_file_args)
        try:
            parts = self._upload_all_parts(upload_id, cloud_file_args)
            self._send_complete_multi_part_upload_request(upload_id, parts, cloud_file_args)
        except Exception as cloudFileUploadException:
            self._send_abort_multi_part_upload_request(upload_id, cloud_file_args)
            raise ConstellationException(
                "Unexpected exception while uploading cloud file", HTTPStatus.INTERNAL_SERVER_ERROR
            ) from cloudFileUploadException
        return self._upload_cloud_file_to_constellation(sharded_key, cloud_file_args)

    @_retry_requests_decorator
    def _send_create_multi_part_upload_request(self, cloud_file_args):
        blobsign_response = self._blobsign("init_multi_upload", cloud_file_args)
        s3_response = requests.post(
            url=blobsign_response["Url"],
        )
        if s3_response.status_code != HTTPStatus.OK:
            raise ConstellationException(
                f"Unexpected non-200 response from s3 endpoint {s3_response}: {s3_response.text}, "
                "failed to init multi-part upload",
                s3_response.status_code,
            )
        parsed_response = xmltodict.parse(s3_response.content)
        return (
            parsed_response["InitiateMultipartUploadResult"]["UploadId"],
            parsed_response["InitiateMultipartUploadResult"]["Key"],
        )

    @_retry_requests_decorator
    def _send_complete_multi_part_upload_request(self, upload_id, parts, cloud_file_args):
        blobsign_response = self._blobsign("complete_multi_upload", cloud_file_args, upload_id=upload_id)
        s3_response = requests.post(url=blobsign_response["Url"], data=parts_to_xml(parts))
        if s3_response.status_code != HTTPStatus.OK:
            raise ConstellationException(
                f"Unexpected non-200 response from s3 endpoint {s3_response}: {s3_response.text}, "
                "failed to complete multi-part upload",
                s3_response.status_code,
            )

    @_retry_requests_decorator
    def _send_abort_multi_part_upload_request(self, upload_id, cloud_file_args):
        blobsign_response = self._blobsign("abort_multi_upload", cloud_file_args)
        s3_response = requests.delete(url=blobsign_response["Url"])
        if s3_response.status_code != HTTPStatus.OK:
            rollbar.report_message(
                "Failed to abort multi part upload. "
                f"UploadID: {upload_id}, Key: {blobsign_response['Key']}, "
                f"S3 response: {s3_response} {s3_response.text}",
                "error",
            )
            raise ConstellationException(
                f"Unexpected non-200 response from s3 {s3_response}: {s3_response.text}, "
                "failed to abort multi part upload",
                s3_response.text,
            )

    def _upload_all_parts(self, upload_id, cloud_file_args):
        current_part_number = 1
        parts = []
        chunk = cloud_file_args.file.read(self._UPLOAD_PART_SIZE)
        while chunk != b"":
            part = self._upload_part(upload_id, current_part_number, chunk, cloud_file_args)
            parts.append(part)
            current_part_number += 1
            chunk = cloud_file_args.file.read(self._UPLOAD_PART_SIZE)
        return parts

    @_retry_requests_decorator
    def _upload_part(self, upload_id, part_number, data_chunk, cloud_file_args):
        blobsign_response = self._blobsign("upload_part", cloud_file_args, part=part_number, upload_id=upload_id)
        s3_response = requests.put(
            url=blobsign_response["Url"],
            data=data_chunk,
        )
        if s3_response.status_code != HTTPStatus.OK:
            raise ConstellationException(
                f"Unexpected non-200 response from s3 endpoint {s3_response}: {s3_response.text}, "
                "failed to upload part",
                s3_response.status_code,
            )
        return {"ETag": s3_response.headers["ETag"], "PartNumber": part_number}

    # Getting history of changes to an object
    def object_log(
        self,
        object_ids=None,
        types=None,
        start_date=None,
        end_date=None,
        max_results=1000,
    ):
        """Find changes on the supplied types after the supplied start_date"""
        payload = {
            "limit": max_results,
            "order": "dsc",
        }
        if object_ids is not None:
            payload["object"] = [{"id": x} for x in object_ids]
        if types is not None:
            payload["Types"] = types
        if start_date is not None:
            payload["startDate"] = start_date
        if end_date is not None:
            payload["endDate"] = end_date
        return self._send_and_validate_request(
            "/obj/objectlog",
            payload,
        )

    # Getting type information
    def get_type(self, type_name):
        """Retrieve all of the fields and their formats for a given type"""
        return self._send_and_validate_request(
            "/obj/type/{}".format(type_name),
            {},
            method="GET",
        )

    # Internal helper functions
    def _auth_header_from_token(self, auth_token):
        """Build a authorization header when an auth token is supplied"""
        return {"Authorization": "Bearer " + str(auth_token)}

    def _send_and_validate_request(
        self,
        path,
        data,
        field_to_retrieve=None,
        method="POST",
        requires_auth_token=True,
        additional_headers=None,
    ):
        """Send a request and validate a request to the given path"""
        if requires_auth_token and not self._auth_token:
            raise ConstellationNotLoggedInException()

        url = self._host + path
        if additional_headers is None:
            additional_headers = {}
        headers = {"Content-type": "application/json", **additional_headers}
        if self._auth_token:
            headers.update(self._auth_header_from_token(self._auth_token))
        if method == "POST":
            response = requests.post(
                url=url,
                data=json.dumps(data),
                headers=headers,
            )
        elif method == "PUT":
            response = requests.put(
                url=url,
                data=json.dumps(data),
                headers=headers,
            )
        elif method == "GET":
            response = requests.get(url=url, headers=headers)
        else:
            raise ConstellationInvalidMethodException(method, path)

        if response.status_code != HTTPStatus.OK:
            raise ConstellationException(
                f"Unexpected non-200 response from constellation endpoint {url}: {response.text}",
                response.status_code,
            )

        if field_to_retrieve:
            field_value = response.json().get(field_to_retrieve)
            if not field_value:
                raise ConstellationInvalidResponseException(
                    f"Could not retrieve field {field_to_retrieve} from request to {path}",
                    response.text,
                )
            return field_value

        return response.json()
