# AUTOGENERATED - modify shared_anyscale_util in root directory to make changes
import dataclasses
from typing import Any, Optional, Type, TypeVar


AWS_ROLE_ARN_RESOURCE_ID_PREFIX = "role/"
_AwsArnT = TypeVar("_AwsArnT", bound="AwsArn")
NoneType = type(None)


@dataclasses.dataclass
class AwsArn:
    partition: str
    service: str
    resource_id: str
    region: Optional[str] = dataclasses.field(default=None)
    account_id: Optional[str] = dataclasses.field(default=None)

    def _validate_fields(self):
        """Enforce typing of each field."""
        if not isinstance(self.partition, str):
            raise ValueError(f"Partition is of type {type(self.partition)}, not str")

        if not isinstance(self.service, str):
            raise ValueError(f"Service is of type {type(self.service)}, not str")

        if not isinstance(self.region, (str, type(None))):
            raise ValueError(
                f"Region is of type {type(self.region)}, not Optional[str]"
            )

        if not isinstance(self.account_id, (str, type(None))):
            raise ValueError(
                "Account ID is of type {}, not Optional[str]".format(
                    type(self.account_id)
                )
            )

        if not isinstance(self.resource_id, str):
            raise ValueError(f"ResourceId is of type {type(self.resource_id)}, not str")

    def _empty_str_to_none(self):
        """Sets empty strings to None."""
        for k, v in dataclasses.asdict(self).items():
            if v == "":
                setattr(self, k, None)

    def __post_init__(self):
        # pydantic.validator("*", pre=True)
        self._empty_str_to_none()

        # Standard pydantic validation
        self._validate_fields()

    @classmethod
    def from_string(cls: Type[_AwsArnT], arn: str) -> _AwsArnT:
        """Convert an AWS ARN string into a Pydantic Model.

        The ARN is a standard format in AWS with
        the following format:
        arn:<partition>:<service>:<region>:<account>:<resource>
        0       1          2         3        4         5

        > arn:aws:apprunner:us-west-2:959243851260:my_app

        Region & Account can be empty i.e:
        > arn:aws:s3:::bucketName
        """
        spl = arn.split(":")
        if len(spl) != 6 or spl[0] != "arn":
            raise ValueError(f"invalid arn: {arn}")
        return cls(
            partition=spl[1],
            service=spl[2],
            region=spl[3],
            account_id=spl[4],
            resource_id=spl[5],
        )

    def to_string(self) -> str:
        return "arn:{p}:{s}:{reg}:{a}:{rid}".format(
            p=self.partition,
            s=self.service,
            reg=self.region or "",
            a=self.account_id or "",
            rid=self.resource_id,
        )


@dataclasses.dataclass
class AwsRoleArn(AwsArn):
    """Role ARNs are of the format arn:aws:iam::<account>:role/<role name>."""

    account_id: str
    region: NoneType

    def _validate_role(self):
        """Validate that a Role has the correct ResourceId."""
        if self.region is not None:
            raise ValueError("Roles do not have a region!")
        if self.account_id is None:
            raise ValueError("No AccountId")
        if not self.resource_id.startswith(AWS_ROLE_ARN_RESOURCE_ID_PREFIX):
            raise ValueError(f"ResourceId {self.resource_id} is not a role")

    def __post_init__(self):
        self.partition = "aws"
        self.service = "iam"

        super().__post_init__()
        self._validate_role()

    @classmethod
    def from_role_name(
        cls: Type[_AwsArnT], account_id: str, role_name: str
    ) -> _AwsArnT:
        return cls(
            partition="aws",
            service="iam",
            account_id=account_id,
            resource_id=f"role/{role_name}",
        )

    def to_role_name(self) -> str:
        """Get the role name from a role ARN.
        If there is a path, this gets the actual name (which is following the last slash).
        Role with path (`/anyscale/`): arn:aws:iam::959243851260:role/anyscale/thomas-test-role
        The Role name *is* `thomas-test-role`
        NOTE: The resource_id for a Role is "role/<role_name>".
        """
        return self.resource_id.split("/")[-1]


def get_dataplane_role_arn(
    account_id: str, cloud_resource_record: Optional[Any],
) -> AwsRoleArn:
    """ Helper method to get instance role from cloud_resource_record. """
    if cloud_resource_record and cloud_resource_record.aws_iam_role_arns:
        return AwsRoleArn.from_string(cloud_resource_record.aws_iam_role_arns[1])

    return AwsRoleArn.from_role_name(
        account_id=account_id, role_name="ray-autoscaler-v1",
    )
