from dataclasses import dataclass
import hashlib
import logging
import shutil
import tarfile
from pathlib import Path
import gzip
import os
from typing import List

logger = logging.getLogger(__name__)

class HashingWriter:
    def __init__(self, base_writer, hash_func=None):
        self.base_writer = base_writer
        self.hash_func = hash_func or hashlib.sha256()

    def write(self, data: bytes):
        self.hash_func.update(data)
        return self.base_writer.write(data)

    def tell(self):
        return self.base_writer.tell()

    def close(self):
        self.base_writer.close()

class MIMETypes:
    mlmodel = "application/x-mlmodel"
    octet_stream = "application/octet-stream"

def get_file_hash(path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(4096):
            h.update(chunk)
    return h.hexdigest()


@dataclass
class LayerStats:
    layer_digest: str
    diff_id: str # will be same as layer_digest if only tar, not targz
    title: str


def tar_filter_fn(input: tarfile.TarInfo) -> tarfile.TarInfo :
    """
    A filter function for modifying file metadata when adding files to a tar archive.

    See also: https://docs.openshift.com/container-platform/4.17/openshift_images/create-images.html#use-uid_create-images

    Args:
        input (tarfile.TarInfo): The file metadata object representing a file being added to the tar archive.

    Returns:
        tarfile.TarInfo: The modified file metadata with the following changes:
            - `uid` set to 0 (root user).
            - `gid` set to 0 (root group).
            - `mode` set to 0o664 (read/write for root owner and root group, read-only for others).
    """
    input.uid = 0
    input.gid = 0
    input.mode = 0o664
    return input


def tarball_from_file(file_path: Path, dest: Path, prefix: str = "/models/") -> LayerStats:
    """
    Creates a tarball from a specified file, storing it in the destination directory with a name based on its checksum.

    Args:
        file_path (Path): The path of the file or directory to be added to the tar archive.
        dest (Path): The destination directory where the tarball will be saved.
        prefix (str): The prefix path for the file_path in the tar. Defaults to `/models/` (KServe ModelCar).

    Returns:
        LayerStats: including the checksum (SHA-256) of the tarball, which is also used as the filename.

    Raises:
        ValueError: If supplied input file path does not exist.
        OSError: If an error occurs during file operations such as renaming or writing.
        tarfile.TarError: If an error occurs during the tar archive creation process.
    """
    if not file_path.exists():
        raise ValueError(f"Input file '{file_path}' does not exist.")
    if not prefix.endswith('/'):
        raise ValueError(f"Supplied prefix {prefix} should end with '/'.")
    
    os.makedirs(dest, exist_ok=True)
    temp_dest = dest / "temp"

    try:
        with open(temp_dest, "wb") as temp_file:
            writer = HashingWriter(temp_file)
            with tarfile.open(fileobj=writer, mode="w") as tar: # type: ignore[call-overload]
                tar.add(file_path, arcname=prefix+file_path.name, filter=tar_filter_fn)
        checksum = writer.hash_func.hexdigest()
        os.rename(temp_dest, dest / checksum)
        return LayerStats(checksum, checksum, file_path.name)
    except tarfile.TarError as e:
        raise tarfile.TarError(f"Error creating tarball: {e}") from e
    except OSError as e:
        raise OSError(f"File operation failed: {e}") from e


def targz_from_file(file_path: Path, dest: Path, prefix: str = "/models/") -> LayerStats:
    """
    Creates a gzipped tarball from the specified file, storing it in the destination directory.
    The tarball's filename is based on the post-compression checksum.

    Args:
        file_path (Path): The path of the file or directory to be compressed.
        dest (Path): The destination directory where the gzipped tarball will be saved. If the directory
                     does not exist, it will be created.
        prefix (str): The prefix path for the file_path in the tar. Defaults to `/models/` (KServe ModelCar).

    Returns:
        LayerStats: containing:
            - The post-compression checksum (SHA-256 of the gzipped tarball).
            - The pre-compression checksum (SHA-256 of the tar content).

    Raises:
        ValueError: If the input file or directory does not exist.
        OSError: If a file operation (e.g., writing, renaming) fails.
        tarfile.TarError: If an error occurs during tarball creation.
    """
    if not file_path.exists():
        raise ValueError(f"Input file '{file_path}' does not exist.")
    if not prefix.endswith('/'):
        raise ValueError(f"Supplied prefix {prefix} should end with '/'.")

    os.makedirs(dest, exist_ok=True)
    temp_dest = dest / "temp"

    try:
        with open(temp_dest, "wb") as temp_file:
            writer = HashingWriter(temp_file)
            with gzip.GzipFile(fileobj=writer, mode="wb", mtime=0, compresslevel=6) as gz: # type: ignore[call-overload]
                inner_writer = HashingWriter(gz)
                with tarfile.open(fileobj=inner_writer, mode="w") as tar: # type: ignore[call-overload]
                    tar.add(file_path, arcname=prefix+file_path.name, filter=tar_filter_fn)
        precompress_checksum = inner_writer.hash_func.hexdigest()
        postcompress_checksum = writer.hash_func.hexdigest()
        os.rename(temp_dest, dest / postcompress_checksum)
        return LayerStats(postcompress_checksum, precompress_checksum, file_path.name)
    except tarfile.TarError as e:
        raise tarfile.TarError(f"Error creating tarball: {e}") from e
    except OSError as e:
        raise OSError(f"File operation failed: {e}") from e


def handle_remove(path: os.PathLike):
    if not isinstance(path, Path):
        path = Path(path)
    if path.is_symlink():
        logger.warning("removing %s which is a symlink", path)
    if path.is_dir():
        shutil.rmtree(path)
    else:
        os.remove(path)


def walk_files(root_path: os.PathLike) -> List[Path]:
    """
    Recursively walks a directory and returns all files as relative paths, skipping any symlinks.
    
    Args:
        root_path: The root directory to walk recursively
        
    Returns:
        List of relative file paths as strings
        
    Raises:
        ValueError: If the provided path doesn't exist or isn't a directory
        OSError: If an error occurs during directory traversal
    """
    if not isinstance(root_path, Path):
        root_path = Path(root_path)
    
    if not root_path.exists():
        raise FileNotFoundError(f"Path '{root_path}' does not exist")
    
    if not root_path.is_dir():
        raise NotADirectoryError(f"Path '{root_path}' is not a directory")
    
    try:
        relative_files = []
        for dirpath, dirnames, filenames in os.walk(str(root_path)):
            # Skip symlink directories by removing them from dirnames
            dirnames[:] = [d for d in dirnames if not os.path.islink(os.path.join(dirpath, d))]
            
            for filename in filenames:
                full_file_path = os.path.join(dirpath, filename)
                # Skip symlink files
                if not os.path.islink(full_file_path):
                    file_path = Path(full_file_path)
                    relative_path = file_path.relative_to(root_path)
                    relative_files.append(relative_path)
        
        return sorted(relative_files)  # Return sorted for consistent ordering
    except OSError as e:
        raise OSError(f"Error walking directory '{root_path}': {e}") from e
