import os
from typing import Any
import requests
import base64
import json
import zipfile


def create_zip_archive(file_paths: list[str], archive_name: str ="archive.zip") -> str:
    """
    Create a ZIP archive from a list of file paths.

    Parameters:
    - file_paths (list of str): List of file paths to include in the ZIP.
    - archive_name (str): Name of the resulting ZIP file.

    Returns:
    - archive file path
    """
    with zipfile.ZipFile(archive_name, 'w') as zipf:
        for i, file_path in enumerate(file_paths):
            if os.path.exists(file_path):
                zipf.write(file_path, arcname=f"{i}_{os.path.basename(file_path)}")
    return archive_name


def needs_auth(func):
    def inner(self, *args, **kw):
        if not self.logged_in:
            self.login()
        try:
            return func(self, *args, **kw)
        except requests.exceptions.HTTPError as e:
            self.login()
            return func(self, *args, **kw)

    return inner


class Config:
    '''
    Configuration object for the connection to Discovery and login for the API access.
    '''
    def __init__(self,
                 HOST: str,
                 CLIENT_ID: str,
                 CLIENT_SECRET: str,
                 USERNAME: str,
                 PASSWORD: str):

        self.HOST = HOST
        self.CLIENT_ID = CLIENT_ID
        self.CLIENT_SECRET = CLIENT_SECRET
        self.USERNAME = USERNAME
        self.PASSWORD = PASSWORD


class Client:
    '''
    Discovery API Client implementation for authentication and file upload.
    '''

    def __init__(self, config: Config):
        '''
        Constructor. Creates a client based on a given configuration.
        
        Parameters:
        - config (Config): Client configuration
        '''
        self.config = config
        self.refresh_token: str | None = None
        self.access_token: str | None = None
        self.logged_in = False


    def _post(self, url: str|None = None, data: Any = None, headers: dict[str, str]|None = None):
        '''
        Generic function to post a query to the API.
    
        Parameters:
          - url (str): API URL.
          - data (Any): the actual payload
          - headers (dict[str, str]): query headers

        Returns:
          - query return value as a json dict
        '''
        if url is None:
            url = f"{self.config.HOST}/api"

        response = requests.request("POST", url, headers=headers, data=data)
        response.raise_for_status()
        json_data = response.json()

        if 'errors' in json_data:
            raise Exception(json_data)
        return json_data


    def login(self):
        '''
        Login to Discovery
        '''
        data = {
            'grant_type': 'password',
            'username': self.config.USERNAME,
            'password': self.config.PASSWORD,
        }
        credentials = f"{self.config.CLIENT_ID}:{self.config.CLIENT_SECRET}"
        encoded_credentials = base64.b64encode(credentials.encode('ascii')).decode('ascii')
        headers = {'Authorization': 'Basic %s' % encoded_credentials}
        data = self._post(f"{self.config.HOST}/oauth/token/", data, headers)
        self.access_token = data['access_token']
        self.refresh_token = data['refresh_token']
        self.logged_in = True
    

    def logout(self):
        '''
        Logout from Discovery
        '''
        self.access_token = None
        self.refresh_token = None
        self.logged_in = False
    

    @needs_auth
    def _create_file(self, filename: str, workbook_uuid: str) -> tuple[str, str] :
        '''
        Creates a new file entry in Discovery and requests a pre-signed S3 URL for uploading the file content.
        Returns the UUID of the file generated by the system and a pre-signed S3 URL that can be used to upload a file.
        
        Parameters:
          - filename (str): name of the file to be created
          - workbook_uuid (str): uuid of the target workbook

        Returns:
          - signed_url (str): A pre-signed S3 URL that can be used to upload a file. The URL is valid for 15 minutes.
          - uuid (str): The UUID of the file generated by the system. 
        '''
        data = {
            "query":
            """mutation CreateFile($input: CreateFileInput!) { 
                createFile(input: $input) { 
                  uuid
                  signedUploadUrl
                }
              }
            """,
            "variables": {
                "input": {
                    "filename": os.path.basename(filename),
                    "workbookUuid": workbook_uuid
                }
            }
        }
        headers = {
              'Authorization': 'Bearer %s' % self.access_token,
              'Content-Type': 'application/json',
          }
        file = self._post(f"{self.config.HOST}/api", headers=headers, data=json.dumps(data))
        return file['data']['createFile']['signedUploadUrl'], file['data']['createFile']['uuid']
    

    def _upload_to_S3(self, signed_url: str, file_path: str) -> None:
        """
        Uploads a file to an S3 pre-signed URL.
        
        Parameters:
          - signed_url (str): A pre-signed S3 URL generated by a call to _create_file
          - file_path (str): 
        """
        with open(file_path, 'rb') as fd:
          headers = {'Slug': os.path.basename(file_path)}
          response = requests.put(signed_url, data=fd, headers=headers)
          response.raise_for_status()
    

    @needs_auth
    def _register_uploaded_file(self, file_uuid: str) -> None:
        """
        Informs Discovery that an S3 file upload has completed. If the upload was successful, Discovery
        will begin processing the file.

        Parameters:
          - file_uuid (str): The uuid of the file to register, see _create_file.
        """
        data = {
            "query":
            """
            mutation RegisterUploadedFile($uuid: String!) {
              registerUploadedFile(uuid: $uuid) {
                uuid
              }
            }
            """,
            "variables": {
                "uuid": file_uuid
            }
          }
        headers = {
              'Authorization': 'Bearer %s' % self.access_token,
              'Content-Type': 'application/json',
          }
        self._post(f"{self.config.HOST}/api", headers=headers, data=json.dumps(data))
        

    @needs_auth
    def upload_file(self, file_path: str, workbook_uuid: str) -> None:
        '''
        Upload a single file to Discovery into the specified Workbook.

        Parameters: 
          - file_path (str): the file to be uploaded
          - workbook_uuid (str): uuid of the target workbook
        '''
        signed_url, file_uuid = self._create_file(os.path.basename(file_path), workbook_uuid)
        self._upload_to_S3(signed_url, file_path)
        self._register_uploaded_file(file_uuid)
    

    @needs_auth
    def upload_files(self, file_paths: list[str], acquisition_name: str, workbook_uuid: str) -> None:
        '''
        Upload a series of files as a zipped archive to Discovery into the specified Workbook.

        Parameters: 
          - file_paths (list of str): the files to be uploaded
          acquisition_name (str): the name of the zipped archive to be uploaded to discovery
          - workbook_uuid (str): uuid of the target workbook
        '''

        if not acquisition_name.endswith(".zip"):
            acquisition_name += ".zip"

        try:
            create_zip_archive(file_paths, acquisition_name)
            signed_url, file_uuid = self._create_file(acquisition_name, workbook_uuid)
            self._upload_to_S3(signed_url, acquisition_name)
            self._register_uploaded_file(file_uuid)
        finally:
            try:
              os.remove(acquisition_name)
            except OSError:
              pass
