from dataclasses import dataclass

from twitch_airflow_components.operators.sagemaker_processing import SageMakerProcessingOperator

from conductor.operators.operator_iface import IConfiguredOperator, TaskWrapper
from conductor.utils.naming import timestamp_name


@dataclass
class BadgerExportOutputs:
    data_s3_url: str
    manifest_s3_url: str


class ConfiguredBadgerExportOperator(IConfiguredOperator):
    def generate_tasks(
        self,
        *,
        task_id: str,
        input_s3_url: str,
        key_field: str,
        dataset_name: str,
        version: str,
        instance_type: str = "ml.m5.4xlarge",
        volume_size_in_gb: int = 128,
    ) -> TaskWrapper[SageMakerProcessingOperator, BadgerExportOutputs]:
        env_config = self.project_resources.env

        env = self.project_resources.env_name
        output_s3_path = self.dag_resources.s3_url_for_path([task_id])
        manifest_s3_path = self.dag_resources.s3_url_for_path(
            [task_id, "manifests", dataset_name, version],
            False,  # disable prefixing to get a consistent path between runs
        )

        # Base configuration assuming Conductor preset values.
        base_config = {
            "ProcessingResources": {
                "ClusterConfig": {
                    "InstanceCount": 1,
                    "InstanceType": instance_type,
                    "VolumeSizeInGB": volume_size_in_gb,
                }
            },
            "ProcessingJobName": timestamp_name(
                f"{task_id}-{self.dag_resources.name}", 63
            ),
            "AppSpecification": {
                "ImageUri": "851170259539.dkr.ecr.us-west-2.amazonaws.com/safety-ml-docker-badger-processor",
                "ContainerEntrypoint": ["processor"],
                "ContainerArguments": [
                    f"-key={key_field}",
                    f"-name={dataset_name}",
                    f"-version={version}",
                    f"-environment={env}",
                    "-timestamp=" + "{{ts}}",
                    f"-manifests3uri={output_s3_path}",
                ],
            },
            "Environment": {
                "ENV": env,
                "GIT_BRANCH": self.project_resources.branch,
                "AWS_DEFAULT_REGION": env_config.default_region,
            },
            "ProcessingInputs": [
                {
                    "InputName": "inputdata",
                    "AppManaged": False,
                    "S3Input": {
                        "S3Uri": input_s3_url,
                        "LocalPath": "/opt/ml/processing/inputdata",
                        "S3DataType": "S3Prefix",
                        "S3InputMode": "File",
                        "S3DataDistributionType": "FullyReplicated",
                        "S3CompressionType": "None",
                    },
                }
            ],
            "ProcessingOutputConfig": {
                "Outputs": [
                    {
                        "OutputName": "outputdata",
                        "S3Output": {
                            "S3Uri": output_s3_path,
                            "LocalPath": "/opt/ml/processing/outputdata",
                            "S3UploadMode": "EndOfJob",
                        },
                    },
                    {
                        "OutputName": "manifest",
                        "S3Output": {
                            "S3Uri": manifest_s3_path,
                            "LocalPath": "/opt/ml/processing/manifest",
                            "S3UploadMode": "EndOfJob",
                        },
                    },
                ]
            },
            "RoleArn": self.project_resources.sagemaker_execution_role(),
        }
        if env_config.vpc is not None:
            network_config = {
                "VpcConfig": {
                    "SecurityGroupIds": env_config.vpc.security_groups,
                    "Subnets": env_config.vpc.subnets,
                }
            }
            base_config["NetworkConfig"] = network_config
        return TaskWrapper[SageMakerProcessingOperator, BadgerExportOutputs](
            SageMakerProcessingOperator(
                config=base_config, task_id=task_id, aws_conn_id=None, dag=self.dag
            ),
            BadgerExportOutputs(
                output_s3_path,
                manifest_s3_path,
            ),
        )
