import boto3
import json
import logging
import os
import pandas
import requests
import uuid

from io import StringIO
from prsn_athena.logger import get_logger
from typing import Any, Dict, List


class AthenaConnection:
    def __init__(
        self,
        database: str,
        workgroup: str,
        output_location: str = None,
        boto3_client=None,
        logger: logging.Logger = None
    ):
        try:
            self._database = database
            self._workgroup = workgroup
            self._output_location = output_location
            self._logger = logger if logger else get_logger('AthenaConnection')
            self._boto3_client = boto3_client if boto3_client else boto3.client('athena')

            if not self._output_location:
                resp = self._boto3_client.get_work_group(
                    WorkGroup=self._workgroup
                )
                wg_cfg = resp.get('WorkGroup', {}).get('Configuration', {})
                output_location = wg_cfg.get('ResultConfiguration', {}).get('OutputLocation')
                if output_location is None:
                    raise Exception(f'No output location found for workgroup {self._workgroup}')
                self._output_location = output_location
        except Exception as e:
            raise Exception(f'Failed to initialize an AthenaConnection with exception: {e}')

    @property
    def database(self) -> str:
        return self._database

    @property
    def workgroup(self) -> str:
        return self._workgroup

    @property
    def output_location(self) -> str:
        return self._output_location

    def execute_query(self, sql: str) -> Dict[str, Any]:
        self._logger.debug(f'Executing sql:\n{sql}')
        try:
            resp = self._boto3_client.start_query_execution(
                QueryString=sql,
                QueryExecutionContext={
                    'Database': self._database
                },
                ResultConfiguration={
                    'OutputLocation': self._output_location
                },
                WorkGroup=self._workgroup
            )

            result = {
                'query_execution_id': resp['QueryExecutionId'],
                'query_execution_status': 'QUEUED',
                'query_execution_start': None,
                'query_execution_end': None,
                'query_results': None
            }
        except Exception as e:
            self._logger.error(f'Error while executing query: {e}')
            raise e

        return result

    def get_query_results(self, query_execution_id: str):
        self._logger.debug(f'Retrieving results for query {query_execution_id}')
        try:
            resp = self._boto3_client.get_query_execution(
                QueryExecutionId=query_execution_id)
            query_execution = resp.get('QueryExecution', {})
            status = query_execution['Status']['State']
            start_time = query_execution['Status'].get('SubmissionDateTime')
            end_time = query_execution['Status'].get('CompletionDateTime')
            result = {
                'query_execution_id': query_execution_id,
                'query_execution_status': status,
                'query_execution_start': None,
                'query_execution_end': None,
                'query_results': None
            }
            if start_time:
                result['query_execution_start'] = start_time.strftime("%Y-%m-%dT%H:%M:%SZ")
            if end_time:
                result['query_execution_end'] = end_time.strftime("%Y-%m-%dT%H:%M:%SZ")
            if status in ['SUCCEEDED']:
                result['query_results'] = self.get_signed_url(
                    query_execution_id=query_execution_id,
                    location=query_execution['ResultConfiguration']['OutputLocation']
                )
            elif status in ['QUEUED', 'RUNNING', 'FAILED', 'CANCELLED']:
                pass
            else:
                raise Exception(f'Unrecognized athena query status: {status}')
            return result
        except Exception as e:
            self._logger.error(f'Error while getting query execution {query_execution_id}: {e}')
            raise e

    def get_signed_url(self, query_execution_id: str, location: str) -> str:
        self._logger.debug(f'Getting a signed URL for query_execution_id: {query_execution_id}')
        try:
            resp = boto3.client('s3').generate_presigned_url(
                'get_object',
                Params={
                    'Bucket': trim_prefix(location, 's3://').split('/', 1)[0],
                    'Key': trim_prefix(location, 's3://').split('/', 1)[1]
                },
                ExpiresIn=900
            )
        except Exception as e:
            self._logger.error(f'Failed to retrieved signed url with error: {e}')
            raise e
        return resp

    def download_results(self, s3_url: str, filepath: str):
        self._logger.debug(f'Downloading file from signed url {s3_url}')
        try:
            with requests.get(s3_url, stream=True) as resp:
                resp.raise_for_status()
                with open(filepath, 'wb') as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        f.write(chunk)
        except Exception as e:
            self._logger.error(f'Failed to download S3 file with exception: {e}')
            raise e

    def download_results_as_dict(self, s3_url: str, filepath: str = None) -> List[Any]:
        self._logger.debug(f'Downloading results as dict')
        try:
            filepath = filepath if filepath else f'/tmp/{uuid.uuid4().hex}.csv'
            self.download_results(s3_url, filepath)
            with open(filepath, 'r') as f:
                df = pandas.read_csv(StringIO(f.read()), dtype=str)
            os.remove(filepath)
            return df.to_dict('records')
        except Exception as e:
            self._logger.error(f'Failed downloading results as dict with exception: {e}')
            raise e

    def to_dict(self) -> Dict[str, str]:
        return {
            'database': self._database,
            'workgroup': self._workgroup,
            'output_location': self._output_location
        }

    def __repr__(self) -> str:
        return json.dumps(self.to_dict())


def trim_prefix(s: str, prefix: str):
    if s.startswith(prefix):
        return s[len(prefix):]
    return s
