# src/chitai/services/utils.py

# Standard library
from __future__ import annotations

import hashlib
from pathlib import Path
import shutil
from typing import TYPE_CHECKING, BinaryIO

if TYPE_CHECKING:
    from hashlib import _Hash

# Third-party libraries
import PIL
import PIL.Image
import aiofiles
import aiofiles.os as aios
from litestar.datastructures import UploadFile


##################################
#  KOReader file hash utilities  #
##################################

# KOReader partial MD5 constants
# These match KOReader's partial MD5 implementation for document identification
# KOReader samples 1024 bytes at specific offsets calculated using 32-bit left shift.
# The shift wrapping behavior (shift & 0x1F) causes i=-1 to produce offset 0.
# Offsets: 0, 1024, 4096, 16384, 65536, 262144, 1048576, ...
KO_STEP = 1024
KO_SAMPLE_SIZE = 1024
KO_INDICES = range(-1, 11)  # -1 to 10 inclusive


def _lshift32(val: int, shift: int) -> int:
    """
    32-bit left shift matching LuaJIT's bit.lshift behavior.

    LuaJIT masks the shift amount to 5 bits (0-31) and performs 32-bit arithmetic.
    This causes negative shifts to wrap: shift=-2 becomes shift=30, and
    1024 << 30 overflows 32 bits to produce 0.
    """
    val &= 0xFFFFFFFF
    shift &= 0x1F
    return (val << shift) & 0xFFFFFFFF


def _get_koreader_offsets() -> list[int]:
    """Get all KOReader sampling offsets."""
    return [_lshift32(KO_STEP, 2 * i) for i in KO_INDICES]


def _partial_md5_from_chunk(
    chunk: bytes,
    hasher: hashlib._Hash,
    offsets: list[int],
    chunk_start: int,
) -> None:
    """
    Update partial MD5 hasher with sampled bytes from a chunk.

    KOReader samples 1024 bytes at specific offsets rather than hashing
    the entire file. This function checks if any sampling offsets fall
    within the current chunk and updates the hasher with those bytes.

    Args:
        chunk: The current chunk of file data.
        hasher: The MD5 hasher to update.
        offsets: List of byte offsets to sample from the file.
        chunk_start: The starting byte position of this chunk in the file.
    """
    chunk_len = len(chunk)
    for offset in offsets:
        if chunk_start <= offset < chunk_start + chunk_len:
            start = offset - chunk_start
            end = min(start + KO_SAMPLE_SIZE, chunk_len)
            hasher.update(chunk[start:end])


async def calculate_koreader_hash(file_path: Path) -> str:
    """
    Calculate KOReader-compatible partial MD5 hash for a file.

    KOReader uses a partial MD5 algorithm that samples 1024 bytes at specific
    offsets rather than hashing the entire file. This provides fast document
    identification for large ebook files.

    The offsets are calculated using 32-bit left shift: 1024 << (2*i) for i from -1 to 10.
    Due to 32-bit overflow, i=-1 produces offset 0:
    0, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, ...

    Args:
        file_path: Path to the file to hash.

    Returns:
        The hexadecimal MD5 hash string.
    """
    hasher = hashlib.md5()
    offsets = _get_koreader_offsets()

    file_pos = 0
    chunk_size = 262144  # 256 KiB

    async with aiofiles.open(file_path, "rb") as f:
        while chunk := await f.read(chunk_size):
            _partial_md5_from_chunk(chunk, hasher, offsets, file_pos)
            file_pos += len(chunk)

    return hasher.hexdigest()


class StreamingHasher:
    """
    Helper class for calculating KOReader hash while streaming file data.

    Allows hash calculation during file writes without needing to re-read
    the file after writing.
    """

    def __init__(self) -> None:
        self.hasher = hashlib.md5()
        self.offsets = _get_koreader_offsets()
        self.position = 0

    def update(self, chunk: bytes) -> None:
        """Update hash with a chunk of data."""
        _partial_md5_from_chunk(chunk, self.hasher, self.offsets, self.position)
        self.position += len(chunk)

    def hexdigest(self) -> str:
        """Return the final hash."""
        return self.hasher.hexdigest()

##################################
#  Filesystem related utilities  #
##################################


class DirectoryDoesNotExist(Exception):
    """Raised when an operation requires a directory that does not exist."""

    ...


async def directory_exists(dir_path: Path | str):
    """
    Check if a directory exists.

    Args:
        dir_path: The directory path to check.

    Returns:
        True if the directory exists, False otherwise.
    """

    return await aios.path.exists(dir_path)


async def create_directory(dir_path: Path | str) -> None:
    """
    Create a directory if it doesn't already exist.

    Prints a message indicating whether the directory was created or already existed.

    Args:
        dir_path: The directory path to create.
    """

    await aios.makedirs(dir_path, exist_ok=True)

async def move_file(src_path: Path, dest_path: Path, create_dirs=True) -> None:
    """
    Move a file from source to destination asynchronously.
    
    Args:
        source_path: Path to the source file
        destination_path: Path to the destination file
    """

    # Create destination directory if needed
    if create_dirs:
        dest_dir = dest_path.parent
        if dest_dir:  # Only create if there's a directory path
            await aios.makedirs(dest_dir, exist_ok=True)

    await aios.rename(src_path, dest_path)

async def move_dir_contents(source_dir: Path | str, target_dir: Path | str) -> None:
    """
    Move all contents from source directory to target directory.

    Creates the target directory if it doesn't exist. Raises an error if the
    source directory does not exist or is not a directory.

    Args:
        source_dir: The directory to move contents from.
        target_dir: The directory to move contents to.

    Raises:
        FileNotFoundError: If source_dir does not exist.
        NotADirectoryError: If source_dir is not a directory.
    """

    source_dir = Path(source_dir)
    target_dir = Path(target_dir)

    if not source_dir.exists():
        raise FileNotFoundError(f"The directory {source_dir} does not exist")
    if not source_dir.is_dir():
        raise NotADirectoryError(f"The path {source_dir} is not a directory")

    if not target_dir.exists():
        target_dir.mkdir(parents=True, exist_ok=True)

    for item in source_dir.iterdir():
        shutil.move(item, target_dir / item.name)


# TODO: Refactor delete_directory to use async I/O for better performance in async contexts.
def delete_directory(directory_path: Path | str):
    """
    Recursively delete a directory and all its contents.

    TODO: Consider changing the implementation to async

    Args:
        directory_path: The directory to delete.
    """

    shutil.rmtree(directory_path)


def cleanup_empty_parent_directories(path: Path, root_path: Path) -> None:
    """
    Remove empty directories from path up to (but not including) root_path.

    Walks up the directory tree from the given path, deleting each empty directory
    until reaching the root_path or a non-empty directory.

    Args:
        path: The starting directory path.
        root_path: The root directory to stop at (not included in deletion).
    """

    # Ensure that the path is a subpath of the root
    if not path.is_relative_to(root_path):
        raise ValueError("Path must be a subdirectory of root path.")

    while path != root_path:
        if path.exists() and is_empty(path):
            delete_directory(path)
        path = path.parent


def is_empty(directory: Path | str):
    """
    Check if a directory is empty.

    Args:
        directory: The directory path to check.

    Returns:
        True if the directory is empty, False otherwise.

    Raises:
        ValueError: If the path does not exist or is not a directory.
    """

    path = Path(directory)

    if not path.exists() or not path.is_dir():
        raise ValueError(f"{directory} is not a valid directory")

    return not any(path.iterdir())


async def delete_file(filepath: Path | str):
    """
    Idempotently delete a file.

    Args:
        filepath: The file path to delete.
    """
    try:
        await aios.remove(filepath)
    except FileNotFoundError:
        pass


async def create_temp_file(data: BinaryIO | bytes):
    """
    Create a temporary file with the given data.

    Args:
        data: Binary data or file-like object to write to the temporary file.

    Returns:
        Path to the created temporary file.
    """

    async with aiofiles.tempfile.NamedTemporaryFile(
        "wb+", delete_on_close=False, delete=False
    ) as file:
        if isinstance(data, bytes):
            await file.write(data)
        else:
            # Handle file-like objects
            data.seek(0)
            file_data = data.read()
            await file.write(file_data)
        file_path = str(file.name)
    return Path(file_path)


async def save_file(data: BinaryIO | bytes, filepath: Path) -> None:
    """
    Save binary data to a file.

    Handles both raw bytes and file-like objects. If data is a file-like object,
    it will be read and the file pointer will be reset to the beginning.

    Args:
        data: Binary data or file-like object to save.
        filepath: The file path where data will be saved.
    """

    async with aiofiles.open(filepath, "wb+") as f:
        if not isinstance(data, bytes):
            await data.seek(0)
            data = await data.read()

        await f.write(data)
        await data.seek(0)


async def save_image(
    img: PIL.Image.Image, filepath: Path, image_type: str = "WEBP"
) -> None:
    """
    Save a PIL Image to a file.

    Args:
        img: The PIL Image object to save.
        filepath: The file path where the image will be saved.
        image_type: The image format to save in (default: "WEBP").
    """

    img.save(filepath, image_type)


##################################
#  Path and filename operations  #
##################################


def get_file_extension(file: Path | str | UploadFile) -> str | None:
    """
    Extract the file extension from a file path.

    Returns the extension in lowercase without the leading dot.

    Args:
        file_path: The file path as a string, Path object, or UploadFile object.

    Returns:
        The file extension in lowercase (e.g., "pdf", "txt"), or None if the path is empty.
    """

    if isinstance(file, str):
        return Path(file).suffix.lower()[1:]

    elif isinstance(file, Path):
        return file.suffix.lower()[1:]

    elif isinstance(file, UploadFile):
        return Path(file.filename).suffix.lower()[1:]
    
    raise ValueError("file object type is not supported")


def get_filename(file: Path | str, ext: bool = True) -> str:
    """
    Extract the filename from a file path.

    Args:
        file_path: The file path as a string or Path object.
        ext: If True (default), returns the full filename with extension.
             If False, returns just the stem (filename without extension).

    Returns:
        The filename, or None if the path is empty.
    """

    if isinstance(file, UploadFile):
        filename = Path(file.filename)
    elif isinstance(file, str):
        filename = Path(Path(file).name)
    elif isinstance(file, Path):
        filename = Path(file.name)
    else:
        raise ValueError("file object type is not supported")
    
    if ext:
        return str(filename)

    return filename.stem


###############################
#  ISBN Validation utilities  #
###############################


def is_valid_isbn(isbn: str) -> bool:
    """
    Validate an ISBN-10 or ISBN-13 number.

    Automatically detects the ISBN format based on length and validates accordingly.

    Args:
        isbn: The ISBN string to validate.

    Returns:
        True if the ISBN is valid, False otherwise.
    """

    try:
        if len(isbn) == 10:
            return is_valid_isbn10(isbn)
        elif len(isbn) == 13:
            return is_valid_isbn13(isbn)
        else:
            return False
    except:
        return False


def is_valid_isbn10(isbn: str) -> bool:
    """
    Validate an ISBN-10 number using its check digit.

    ISBN-10 uses a weighted sum where each digit is multiplied by its position
    (10 down to 1) and the result modulo 11 determines the check digit.

    Args:
        isbn: The 10-character ISBN string to validate.

    Returns:
        True if the ISBN-10 is valid, False otherwise.
    """
    sum = 0
    for i in range(9):
        sum += int(isbn[i]) * (10 - i)

    check_digit = 11 - (sum % 11)
    return str(check_digit) == isbn[-1] or (check_digit == 10 and isbn[-1] in "Xx")


def is_valid_isbn13(isbn: str) -> bool:
    """
    Validate an ISBN-13 number using its check digit.

    ISBN-13 uses a weighted sum where digits alternate between multipliers of 1 and 3,
    and the result modulo 10 determines the check digit.

    Args:
        isbn: The 13-character ISBN string to validate.

    Returns:
        True if the ISBN-13 is valid, False otherwise.
    """
    sum = 0
    for i in range(12):
        sum += int(isbn[i]) * (1 if i % 2 == 0 else 3)

    check_digit = 10 - (sum % 10)
    return str(check_digit % 10) == isbn[-1]
