import os
from datetime import datetime

from pyspark import SparkConf, SparkContext, RDD


class SparkTask:
    task_name: str
    target_day: datetime
    target_day_str: str

    def __init__(self, app_name: str, day: str):
        """

        Args:
            app_name: task name
            day: date string in the format of '%Y-%m-%d'
        """
        self.task_name = app_name
        self.target_day = datetime.strptime(day, '%Y-%m-%d')
        self.target_day_str = day
        self.sc = SparkContext(conf=SparkConf().setAppName(app_name))

    def files(self) -> RDD:
        """
        Notes: This method is an abstract method, sub-class should override it.

        Returns: Spark RDD
        """
        raise NotImplemented("SparkTask is an abstract class, use it's subclass")


class SparkTaskLocalFile(SparkTask):
    path_pattern: str

    def __init__(self, path_pattern: str, **kwargs):
        """

        Args:
            path_pattern: path pattern to the local file contains could contains date variables (%Y, %m, %d)
            **kwargs: see `SparkTask`.`init`
        """
        super().__init__(**kwargs)
        self.path_pattern = path_pattern

    def files(self) -> RDD:
        return self.sc.textFile(self.target_day.strftime(self.path_pattern))


# see https://stackoverflow.com/questions/32155617/connect-to-s3-data-from-pyspark
# see https://stackoverflow.com/questions/46740670/no-filesystem-for-scheme-s3-with-pyspark
# see https://stackoverflow.com/questions/30851244/spark-read-file-from-s3-using-sc-textfile-s3n
class SparkTaskAmazonS3(SparkTask):
    bucket: str
    region: str
    prefix_pattern: str
    aws_access_key_id: str
    aws_secret_access_key: str
    scheme: str

    def __init__(self, bucket: str, region: str, access_key_id: str, secret_access_key: str, path_pattern: str, scheme: str = "s3n", **kwargs):
        """

        Args:
            bucket: S3 bucket
            region: S3 region
            access_key_id: fs.s3.awsAccessKeyId
            secret_access_key:
            path_pattern: path pattern of s3 objects contains could contains date variables (%Y, %m, %d)
            scheme: defaults to 's3n'
            **kwargs: see `SparkTask`.`init`
        """
        super().__init__(**kwargs)
        self.bucket = bucket
        self.region = region
        self.prefix_pattern = path_pattern
        self.aws_access_key_id = access_key_id
        self.aws_secret_access_key = secret_access_key
        self.scheme = scheme

        hc = self.sc._jsc.hadoopConfiguration()
        hc.set(f'fs.{self.scheme}.impl', 'org.apache.hadoop.fs.s3native.NativeS3FileSystem')
        hc.set(f'fs.{self.scheme}.awsAccessKeyId', self.aws_access_key_id)
        hc.set(f'fs.{self.scheme}.awsSecretAccessKey', self.aws_secret_access_key)

    def files(self) -> RDD:
        prefix = self.target_day.strftime(self.prefix_pattern)
        return self.sc.textFile(f"{self.scheme}://{self.bucket}/{prefix}")