import csv
import logging
import time
from datetime import datetime
from typing import Any, Optional, Union

import boto3
from typeguard import typechecked

from psi_athena import get_logger, log_it

DEFAULT_ATHENA_DATA_CATALOG = "AwsDataCatalog"
DEFAULT_ATHENA_DATABASE = "default"
DEFAULT_ATHENA_WORKGROUP = "primary"
DEFAULT_AWS_REGION = "us-east-1"
DEFAULT_AWS_SESSION_NAME = "psi-athena-session"
DEFAULT_MAX_POLL_DURATION_SECONDS = 300
DEFAULT_POLL_WAIT_SECONDS = 5


@typechecked
class Athena:

    def __init__(
        self,
        data_catalog: Optional[str] = None,
        workgroup: Optional[str] = None,
        database: Optional[str] = None,
        max_poll_duration_seconds: Optional[int] = None,
        poll_wait_seconds: Optional[int] = None,
        aws_session: Optional[Any] = None,
        logger: Optional[logging.Logger] = None,
    ):
        self.data_catalog = data_catalog or DEFAULT_ATHENA_DATA_CATALOG
        self.workgroup = workgroup or DEFAULT_ATHENA_WORKGROUP
        self.database = database or DEFAULT_ATHENA_DATABASE
        self.max_poll_duration_seconds = (
            max_poll_duration_seconds or DEFAULT_MAX_POLL_DURATION_SECONDS
        )
        self.poll_wait_seconds = poll_wait_seconds or DEFAULT_POLL_WAIT_SECONDS
        self.__aws_session = aws_session or boto3.Session()
        self.__athena_client = None
        self.__logger = logger or get_logger()
        self.assume_role()

    def assume_role(
        self,
        role_arn: Optional[str] = None,
        region: Optional[str] = None,
        session_name: Optional[str] = None,
    ):
        if role_arn:
            region = region or DEFAULT_AWS_REGION
            session_name = session_name or DEFAULT_AWS_SESSION_NAME
            self.__logger.info(f"Assuming role {role_arn} in region {region}")
            stsclient = self.__aws_session.client("sts")
            stsresponse = stsclient.assume_role(RoleArn=role_arn, RoleSessionName=session_name)

            self.__logger.debug(f"Initializing Athena client from role: {role_arn}")
            self.__athena_client = self.__aws_session.client(
                "athena",
                region_name=region,
                aws_access_key_id=stsresponse["Credentials"]["AccessKeyId"],
                aws_secret_access_key=stsresponse["Credentials"]["SecretAccessKey"],
                aws_session_token=stsresponse["Credentials"]["SessionToken"],
            )
        else:
            self.__logger.debug("Initializing Athena client from existing AWS session")
            self.__athena_client = self.__aws_session.client("athena")

    @log_it
    def run_athena_query(
        self,
        query: str,
        data_catalog: Optional[str] = None,
        workgroup: Optional[str] = None,
        database: Optional[str] = None,
        table: Optional[str] = None,
        download_results: bool = True,
        reload_partitions: bool = False,
    ) -> Union[dict, list]:
        if reload_partitions:
            self.reload_partitions(table)

        query_execution_id = self.start_query(
            query=query, data_catalog=data_catalog, workgroup=workgroup, database=database
        )
        query_results = self.poll_for_query_results(query_execution_id)

        if download_results:
            return self.download_query_results_as_dict(query_results)
        return query_results

    def start_query(
        self,
        query: str,
        data_catalog: Optional[str] = None,
        workgroup: Optional[str] = None,
        database: Optional[str] = None,
    ) -> str:
        data_catalog = data_catalog or self.data_catalog
        workgroup = workgroup or self.workgroup
        database = database or self.database
        self.__logger.info(f"Starting Athena query: {query}")
        try:
            query_execution_id = self.__athena_client.start_query_execution(
                QueryString=query,
                QueryExecutionContext={"Database": database, "Catalog": data_catalog},
                WorkGroup=workgroup,
            )["QueryExecutionId"]
        except Exception as e:
            self.__logger.error(f"Failed to start Athena query execution: {e}")
            raise e
        self.__logger.info(f"Athena query execution id: {query_execution_id}")
        return query_execution_id

    def poll_for_query_results(self, query_execution_id: str) -> dict:
        t0 = datetime.now()
        query_execution_status = "QUEUED"
        while (datetime.now() - t0).total_seconds() < self.max_poll_duration_seconds:
            try:
                query_results = self.__athena_client.get_query_execution(
                    QueryExecutionId=query_execution_id
                )
                query_execution_status = query_results["QueryExecution"]["Status"]["State"]
            except Exception as e:
                self.__logger.error(f"Failed to retrieve Athena query execution status: {e}")
                raise e

            if query_execution_status in ["SUCCEEDED", "FAILED", "CANCELLED"]:
                break
            time.sleep(self.poll_wait_seconds)
        assert (
            query_execution_status == "SUCCEEDED"
        ), f"Athena query execution did not succeed, polling exited with status: {query_execution_status}"
        return query_results

    def download_query_results_as_dict(self, query_results: dict) -> list[dict]:
        data = self.download_query_results(query_results)
        try:
            csvstr = data.decode("utf-8")
            parsed = list(csv.DictReader(csvstr.split("\n")))
            for i in range(len(parsed)):
                for k, v in parsed[i].items():
                    assert k and isinstance(
                        k, str
                    ), f"Invalid CSV column (likely that extra columns were found): {k}"
                    if len(v) == 0:
                        # Parse empty strings as None
                        parsed[i][k] = None
            return parsed
        except Exception as e:
            self.__logger.error(f"Failed to parse query results into a dict: {e}")
            raise e

    def download_query_results(self, query_results: dict) -> bytes:
        s3_path = (
            query_results.get("QueryExecution", {})
            .get("ResultConfiguration", {})
            .get("OutputLocation")
        )
        try:
            client = self.__aws_session.client("s3")
            self.__logger.debug(f"Fetching object from S3: {s3_path}")
            obj = client.get_object(
                Bucket=s3_path[len("s3://") :].split("/", 1)[0],
                Key=s3_path[len("s3://") :].split("/", 1)[1],
            )
            return obj["Body"].read()
        except Exception as e:
            self.__logger.error(f"Failed to download results with error: {e}")
            raise e

    def reload_partitions(self, table: str):
        assert not table.endswith("_vw"), f"Cannot reload partitions on a view"
        query = f"MSCK REPAIR TABLE `{table}`"
        query_result = self.run_athena_query(
            query, table=table, download_results=False, reload_partitions=False
        )
        assert query_result["QueryExecution"]["Status"]["State"] == "SUCCEEDED"
