from typing import Dict, Any, Optional
from pydantic import BaseModel, Field, model_validator
from enum import Enum


class StorageType(str, Enum):
    S3 = "s3"
    Azure_BLOB = "azure_blob"
    GCS = "gcs"
    LOCAL = "local"


class CredentialsType(str, Enum):
    ACCESS_KEY = "access_key"
    IAM_ROLE = "iam_role"
    ENV_VARS = "env_vars"
    PROFILE = "profile"
    ACCOUNT_KEY = "account_key"
    SERVICE_ACCOUNT = f"service_account"


class StorageCredentials(BaseModel):
    type: CredentialsType
    # AWS credentials
    aws_access_key_id: Optional[str] = None
    aws_secret_access_key: Optional[str] = None
    # Azure credentials
    account_key: Optional[str] = None
    # GCS credentials
    service_account_key_path: Optional[str] = None

    @model_validator(mode="before")
    def validate_credentials(cls, values):
        cred_type = values.get("type")
        if cred_type == CredentialsType.ACCESS_KEY:
            if not (values.get("aws_access_key_id") and values.get("aws_secret_access_key")):
                raise ValueError("aws_access_key and aws_secret_access_key required for access_key type")
        elif cred_type == CredentialsType.ACCOUNT_KEY:
            if not values.get("account_key"):
                raise ValueError("account_key required for account_key type")
        elif cred_type == CredentialsType.SERVICE_ACCOUNT:
            if not values.get("service_account_key_path"):
                raise ValueError("service_account_key_path required for service_account type")

        return values


class StorageConnection(BaseModel):
    # AWS S3 connection
    region: Optional[str] = None
    endpoint_url: Optional[str] = None
    # Azure connection
    account_name: Optional[str] = None
    endpoint_suffix: Optional[str] = None
    # GCS connection
    project_id: Optional[str] = None


class StorageConfig(BaseModel):
    type: StorageType
    connection: Optional[StorageConnection] = None
    credentials: Optional[StorageCredentials] = None
    config: Dict[str, Any] = Field(default_factory=dict)


class StorageConfigs(BaseModel):
    default: str = Field(description="Default storage reference name")
    storages: Dict[str, StorageConfig] = Field(description="Storage configurations by name")

    @model_validator(mode="after")
    def validate_default_exists(self):
        if self.default not in self.storages:
            raise ValueError(f"Default storage '{self.default}' not found in storages")
        return self