import boto3
import re
from botocore.exceptions import ClientError, NoCredentialsError
from urllib.parse import urlparse
from pathlib import Path
from io import IOBase
from typing import Union, List, Optional
from emissary_ml_toolkit.storage.base import BaseStorage
from emissary_ml_toolkit.config.storage import CredentialsType
from emissary_ml_toolkit.utils.exceptions import StorageError


class S3Storage(BaseStorage):
    def _setup_client(self) -> None:
        try:
            client_config = {}
            if self.config.connection.region:
                client_config["region_name"] = self.config.connection.region
            if self.config.connection.endpoint_url:
                client_config["endpoint_url"] = self.config.connection.endpoint_url

            # TODO: Handle multiple credentials type
            credentials = self.config.credentials

            if credentials.type == CredentialsType.ACCESS_KEY:
                client_config.update({
                    "aws_access_key_id": credentials.aws_access_key_id,
                    "aws_secret_access_key": credentials.aws_secret_access_key,
                })

            self.connection = boto3.client('s3', **client_config)

        except NoCredentialsError:
            raise StorageError("AWS credentials not found.")
        except Exception as e:
            raise StorageError(f"Failed to setup S3 client: {str(e)}")

    @staticmethod
    def _parse_s3_path(storage_path: str) -> tuple:
        """Parse S3 path into bucket and object key"""
        s3_url = storage_path.strip()
        if s3_url.startswith("s3://"):
            path = s3_url[5:]  # Remove 's3://'
            if '/' not in path:
                raise StorageError("Invalid S3 URL: missing key component")

            bucket, key = path.split('/', 1)
            if not bucket or not key:
                raise StorageError("Invalid S3 URL: empty bucket or key")
            return bucket, key
            # Parse HTTP/HTTPS URLs
        parsed = urlparse(s3_url)

        if not parsed.netloc:
            raise StorageError("Invalid S3 URL: no hostname found")

        # Virtual hosted-style URL: https://bucket-name.s3.region.amazonaws.com/key
        virtual_hosted_pattern = r'^([^.]+)\.s3\.([^.]+)\.amazonaws\.com$'
        match = re.match(virtual_hosted_pattern, parsed.netloc)

        if match:
            bucket_name = match.group(1)
            # Remove leading slash from path to get the key
            key = parsed.path.lstrip('/')

            if not key:
                raise StorageError("Invalid S3 URL: missing key component")

            return bucket_name, key

        # Path-style URL: https://s3.region.amazonaws.com/bucket-name/key
        path_style_pattern = r'^s3\.([^.]+)\.amazonaws\.com$'
        match = re.match(path_style_pattern, parsed.netloc)

        if match:
            path_parts = parsed.path.strip('/').split('/', 1)

            if len(path_parts) < 2:
                raise StorageError("Invalid S3 URL: missing bucket or key component")

            bucket_name, key = path_parts

            if not bucket_name or not key:
                raise StorageError("Invalid S3 URL: empty bucket or key")

            return bucket_name, key

        # Check for other S3 endpoint patterns (like s3.amazonaws.com)
        if 's3' in parsed.netloc and 'amazonaws.com' in parsed.netloc:
            # Generic S3 path-style parsing
            path_parts = parsed.path.strip('/').split('/', 1)

            if len(path_parts) >= 2:
                bucket_name, key = path_parts
                if bucket_name and key:
                    return bucket_name, key

        raise StorageError(f"Invalid S3 URL format: {s3_url}")


    def download_file(self, storage_path: str, local_path: Optional[Union[str, Path]] = None) -> Path:
        bucket, key = self._parse_s3_path(storage_path)
        if not local_path:
            local_path = storage_path.split('/')[-1]
        local_path = Path(local_path)
        local_path.parent.mkdir(parents=True, exist_ok=True)

        try:
            self.connection.download_file(bucket, key, str(local_path))
            return local_path
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchKey':
                raise StorageError(f"Object not found: {storage_path}")
            if error_code == 'NoSuchBucket':
                raise StorageError(f"Bucket not found: {bucket}")
            raise StorageError(f"Failed to download file from S3: {str(e)}")

    def upload_file(self, source: Union[str, Path, IOBase, bytes], storage_path: Union[str, Path]):
        bucket, key = self._parse_s3_path(storage_path)
        try:
            if isinstance(source, (str, Path)):
                local_path = Path(source)
                if not local_path.exists():
                    raise StorageError(f"Local file not found: {local_path}")
                self.connection.upload_file(local_path, bucket, key)
                return
            if isinstance(source, bytes):
                self.connection.put_object(Bucket=bucket, Key=key, Body=source)
                return
            if isinstance(source, IOBase):
                if hasattr(source, 'getvalue'):
                    body = source.getvalue()
                else:
                    body = source.read()
                self.connection.put_object(Bucket=bucket, Key=key, Body=body)
                return
            raise StorageError(f"Unsupported source type: {type(source)}")
        except ClientError as e:
            raise StorageError(f"Failed to upload file to S3: {str(e)}")

    def list_objects(self, prefix: str = "", bucket: Optional[str] = None) -> List[str]:
        if not bucket:
            raise StorageError("AWS S3 needs bucket to list objects")

        try:
            objects = []
            paginator = self.connection.get_paginator('list_objects_v2')

            for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
                if 'Contents' in page:
                    for obj in page['Contents']:
                        objects.append(obj['Key'])

            return objects
        except ClientError as e:
            raise StorageError(f"Failed to list objects: {str(e)}")

    def exists(self, storage_path: str) -> bool:
        bucket, key = self._parse_s3_path(storage_path)
        try:
            self.connection.head_object(Bucket=bucket, Key=key)
            return True
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchKey':
                return False
            else:
                raise StorageError(f"Failed to check if object exists: {str(e)}")
