"""
S3 client for RunPod network volume operations.

Provides:
- File uploads with progress tracking
- Log file tailing via byte-range requests
- Model download
- Volume inventory listing
"""

import hashlib
import time
from pathlib import Path
from typing import Any, Callable

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TaskID,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)

from .config import RunPodConfig, get_config
from .errors import S3Error, S3ObjectNotFoundError

console = Console()


class S3Client:
    """
    S3 client for RunPod network volume operations.

    Handles all S3 interactions with proper retry logic and error handling.
    """

    def __init__(self, config: RunPodConfig | None = None):
        """
        Initialize S3 client.

        Args:
            config: RunPodConfig instance (uses global config if None)
        """
        self.config = config or get_config()

        # Initialize boto3 S3 client
        self.s3_client = boto3.client(
            "s3",
            aws_access_key_id=self.config.runpod_s3_access_key,
            aws_secret_access_key=self.config.runpod_s3_secret,
            region_name=self.config.runpod_s3_region,
            endpoint_url=self.config.runpod_s3_endpoint,
            config=Config(
                signature_version="s3v4",
                s3={"addressing_style": "path"},
                retries={
                    "max_attempts": self.config.max_retries,
                    "mode": "adaptive",
                },
                connect_timeout=30,
                read_timeout=self.config.api_timeout,
            ),
        )

        self.bucket = self.config.get_s3_bucket()

    def compute_md5(self, file_path: Path) -> str:
        """
        Compute MD5 checksum of local file.

        Args:
            file_path: Path to local file

        Returns:
            MD5 hex digest

        Raises:
            S3Error: If file cannot be read
        """
        try:
            md5 = hashlib.md5()
            with open(file_path, "rb") as f:
                for chunk in iter(lambda: f.read(8192), b""):
                    md5.update(chunk)
            return md5.hexdigest()
        except OSError as e:
            raise S3Error(f"Failed to read file {file_path}: {e}") from e

    def object_exists(self, s3_key: str) -> bool:
        """
        Check if S3 object exists.

        Args:
            s3_key: S3 object key

        Returns:
            True if object exists, False otherwise
        """
        try:
            self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return False
            raise S3Error(f"Failed to check object existence: {e}", self.bucket, s3_key) from e

    def get_object_size(self, s3_key: str) -> int:
        """
        Get size of S3 object in bytes.

        Args:
            s3_key: S3 object key

        Returns:
            Object size in bytes

        Raises:
            S3ObjectNotFoundError: If object doesn't exist
            S3Error: On other S3 errors
        """
        try:
            response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            return response["ContentLength"]
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                raise S3ObjectNotFoundError(self.bucket, s3_key)
            raise S3Error(f"Failed to get object size: {e}", self.bucket, s3_key) from e

    def needs_upload(self, local_file: Path, s3_key: str) -> tuple[bool, str]:
        """
        Check if file needs upload based on MD5 checksum.

        Args:
            local_file: Path to local file
            s3_key: Destination S3 key

        Returns:
            Tuple of (needs_upload, reason)
        """
        local_md5 = self.compute_md5(local_file)

        try:
            response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            remote_etag = response["ETag"].strip('"')

            if local_md5 == remote_etag:
                return False, "up-to-date (checksum match)"
            else:
                return True, "changed (checksum mismatch)"
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return True, "new file (doesn't exist on S3)"
            else:
                return True, f"error checking: {e}"

    def upload_binary(
        self,
        local_file: Path,
        s3_key: str,
        force: bool = False,
        progress_callback: Callable[[int], None] | None = None,
    ) -> bool:
        """
        Upload binary file with progress tracking.

        Args:
            local_file: Path to local file
            s3_key: Destination S3 key
            force: Skip checksum check and force upload
            progress_callback: Optional callback(bytes_transferred)

        Returns:
            True if upload succeeded, False otherwise

        Raises:
            S3Error: On upload failure
        """
        if not local_file.exists():
            raise S3Error(f"Local file not found: {local_file}")

        file_size = local_file.stat().st_size

        # Check if upload needed (unless forced)
        if not force:
            needs_upload, reason = self.needs_upload(local_file, s3_key)
            if not needs_upload:
                console.print(f"[green]✓ {s3_key}[/green]: {reason}")
                return True
            console.print(f"[cyan]↻ {s3_key}[/cyan]: {reason}")

        try:
            # Upload with optional progress callback
            self.s3_client.upload_file(
                str(local_file),
                self.bucket,
                s3_key,
                Callback=progress_callback,
                ExtraArgs={"ContentType": "application/octet-stream"},
            )

            # Verify upload
            response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            remote_size = response["ContentLength"]

            if remote_size == file_size:
                console.print(f"[green]✓ Uploaded: {s3_key} ({file_size:,} bytes)[/green]")
                return True
            else:
                # Size mismatch - delete and raise error
                console.print(
                    f"[red]✗ Size mismatch: local={file_size}, remote={remote_size}[/red]"
                )
                self.s3_client.delete_object(Bucket=self.bucket, Key=s3_key)
                raise S3Error(
                    f"Upload verification failed: size mismatch", self.bucket, s3_key
                )

        except ClientError as e:
            raise S3Error(f"Upload failed: {e}", self.bucket, s3_key) from e

    def tail_log_file(
        self, s3_key: str, start_byte: int = 0, max_bytes: int = 1024 * 1024
    ) -> tuple[bytes, int]:
        """
        Tail log file using byte-range requests.

        Args:
            s3_key: S3 key for log file
            start_byte: Starting byte position
            max_bytes: Maximum bytes to read (default 1MB)

        Returns:
            Tuple of (log_content, new_position)

        Raises:
            S3ObjectNotFoundError: If log file doesn't exist
            S3Error: On other S3 errors
        """
        try:
            # Get current file size
            file_size = self.get_object_size(s3_key)

            # If file hasn't grown, return empty
            if file_size <= start_byte:
                return b"", start_byte

            # Calculate byte range to fetch
            end_byte = min(start_byte + max_bytes - 1, file_size - 1)

            # Fetch byte range
            response = self.s3_client.get_object(
                Bucket=self.bucket, Key=s3_key, Range=f"bytes={start_byte}-{end_byte}"
            )

            content = response["Body"].read()
            new_position = start_byte + len(content)

            return content, new_position

        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchKey":
                raise S3ObjectNotFoundError(self.bucket, s3_key)
            raise S3Error(f"Failed to tail log file: {e}", self.bucket, s3_key) from e

    def download_results(self, s3_prefix: str, local_dir: Path) -> list[Path]:
        """
        Download all files with given S3 prefix to local directory.

        Args:
            s3_prefix: S3 prefix (directory) to download
            local_dir: Local destination directory

        Returns:
            List of downloaded file paths

        Raises:
            S3Error: On download failure
        """
        local_dir.mkdir(parents=True, exist_ok=True)
        downloaded_files = []

        try:
            # List all objects with prefix
            paginator = self.s3_client.get_paginator("list_objects_v2")
            pages = paginator.paginate(Bucket=self.bucket, Prefix=s3_prefix)

            objects = []
            for page in pages:
                objects.extend(page.get("Contents", []))

            if not objects:
                console.print(f"[yellow]⚠ No files found with prefix: {s3_prefix}[/yellow]")
                return downloaded_files

            console.print(f"Found {len(objects)} file(s) to download")

            # Download each object with progress
            with Progress(
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                DownloadColumn(),
                TransferSpeedColumn(),
                TimeRemainingColumn(),
                console=console,
            ) as progress:
                for obj in objects:
                    s3_key = obj["Key"]
                    file_size = obj["Size"]

                    # Skip directory markers
                    if s3_key.endswith("/"):
                        continue

                    # Determine local path (preserve directory structure)
                    rel_path = Path(s3_key).relative_to(s3_prefix)
                    local_file = local_dir / rel_path
                    local_file.parent.mkdir(parents=True, exist_ok=True)

                    # Download with progress
                    task = progress.add_task(f"Downloading {rel_path}", total=file_size)

                    def callback(bytes_transferred):
                        progress.update(task, completed=bytes_transferred)

                    self.s3_client.download_file(
                        self.bucket, s3_key, str(local_file), Callback=callback
                    )

                    downloaded_files.append(local_file)
                    console.print(f"[green]✓ Downloaded: {rel_path}[/green]")

            return downloaded_files

        except ClientError as e:
            raise S3Error(f"Download failed: {e}", self.bucket, s3_prefix) from e

    def list_binaries(self) -> list[dict[str, Any]]:
        """
        List all binaries in binaries/ directory.

        Returns:
            List of dicts with keys: name, size, last_modified
        """
        return self._list_prefix("binaries/")

    def list_models(self) -> list[dict[str, Any]]:
        """
        List all models in models/ directory.

        Returns:
            List of dicts with keys: name, size, last_modified
        """
        models_prefix = self.config.get_models_s3_prefix()
        return self._list_prefix(models_prefix)

    def _list_prefix(self, prefix: str) -> list[dict[str, Any]]:
        """
        List all objects with given prefix.

        Args:
            prefix: S3 prefix to list

        Returns:
            List of object metadata dicts
        """
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)

            objects = []
            for obj in response.get("Contents", []):
                # Skip directory markers
                if obj["Key"].endswith("/"):
                    continue

                objects.append(
                    {
                        "name": obj["Key"],
                        "size": obj["Size"],
                        "last_modified": obj["LastModified"],
                    }
                )

            return objects

        except ClientError as e:
            raise S3Error(f"Failed to list objects: {e}", self.bucket, prefix) from e

    def create_directory(self, s3_key: str) -> bool:
        """
        Create directory marker in S3.

        Args:
            s3_key: Directory path (should end with /)

        Returns:
            True if created or already exists
        """
        if not s3_key.endswith("/"):
            s3_key += "/"

        try:
            # Check if already exists
            if self.object_exists(s3_key):
                return True

            # Create empty directory marker
            self.s3_client.put_object(Bucket=self.bucket, Key=s3_key)
            console.print(f"[green]✓ Created directory: {s3_key}[/green]")
            return True

        except ClientError as e:
            console.print(f"[yellow]⚠ Failed to create directory {s3_key}: {e}[/yellow]")
            return False