# src/chitai/services/book.py

# Standard library
from __future__ import annotations
from collections import defaultdict
import mimetypes

from io import BytesIO
from pathlib import Path
import uuid
import zipfile
from typing import TYPE_CHECKING, Any, AsyncIterator

# Third-party libraries
from advanced_alchemy.extensions.litestar import service
from advanced_alchemy.service import (
    SQLAlchemyAsyncRepositoryService,
    ModelDictT,
    is_dict,
    schema_dump
)
from advanced_alchemy.repository import SQLAlchemyAsyncRepository
from advanced_alchemy.filters import CollectionFilter
from sqlalchemy import inspect
from litestar.response import File
from litestar.datastructures import UploadFile
import aiofiles
from aiofiles import os as aios
from PIL import Image

# Local imports
from chitai.config import settings
from chitai.database.models import (
    Book,
    Author,
    Tag,
    Publisher,
    BookSeries,
    FileMetadata,
    Identifier,
    Library,
)
from chitai.schemas.book import BooksCreateFromFiles
from chitai.services.filesystem_library import BookPathGenerator
from chitai.services.metadata_extractor import Extractor as MetadataExtractor
from chitai.services.utils import (
    calculate_koreader_hash,
    cleanup_empty_parent_directories,
    delete_file,
    move_dir_contents,
    move_file,
    save_image,
    StreamingHasher,
)


class BookService(SQLAlchemyAsyncRepositoryService[Book]):
    """Book service for managing book operations."""

    class Repo(SQLAlchemyAsyncRepository[Book]):
        """Book repository."""

        model_type = Book

    repository_type = Repo

    

    async def create_book(self, data: ModelDictT[Book], library: Library, **kwargs) -> Book:
        """
        Create a new book entity.

        Orchestrates the full book creation pipeline: extracting metadata from files,
        saving cover images, and storing book files to the filesystem.

        Args:
            data: Book data as a dictionary or model.
            library: The library the book belongs to.
            *args: Additional positional arguments passed to parent create.
            **kwargs: Additional keyword arguments passed to parent create.

        Returns:
            The created Book entity.
        """
        data = schema_dump(data)
        await self._parse_metadata_from_files(data)

        await self._save_cover_image(data)
        await self._save_book_files(library, data)

        return await super().create(data, **kwargs)

    async def create_many_from_files(
        self, data: BooksCreateFromFiles, library: Library, **kwargs
    ) -> list[Book]:
        """
        Create multiple books from uploaded files.

        Groups files by their parent directory to organize books. Files in the root
        directory are treated as separate individual books.

        Args:
            data: Container with list of uploaded files.
            library: The library the books belong to.
            *args: Additional positional arguments passed to create.
            **kwargs: Additional keyword arguments passed to create.

        Returns:
            List of created Book entities.
        """

        if not data.files:
            raise ValueError("Must upload at least one file")

        # Group book files if they are within the same nested directory
        books = defaultdict(list)
        for file in data.files:
            filepath = Path(file.filename)
            # Books within the root directory should be treated as separate books
            if len(filepath.parent.parts) > 1:
                books[filepath.parent].append(file)
            else:
                books[filepath].append(file)

        return [
            await self.create_book(
                {"files": [file for file in files], "library_id": library.id},
                library,
                **kwargs,
            )
            for files in books.values()
        ]
    
    async def create_many_from_existing_files(
        self,
        file_paths: list[Path],
        consume_path: Path,
        library: Library,
        **kwargs
    ) -> list[Book]:
        
        # Group up files if they are in the same leaf directory

        books = []
        file_groups: dict[Path, list[Path]] = defaultdict(list)


        for file_path in file_paths:
            rel_path = file_path.relative_to(consume_path)
            parent_rel = rel_path.parent

            # Files at root level should each be in their own group
            # Use the file's relative path as the group key instead of parent
            if parent_rel == Path("."):
                group_key = rel_path  # Use the file itself as the key
            else:
                group_key = parent_rel  # Use parent directory as key
        
            # Add to appropriate group
            file_groups[group_key].append(file_path)


        # For each grouping
        for group, files in file_groups.items():
            data: dict[str, Any] = {'files': files}
            await self._parse_metadata_from_files(data, root_path=consume_path)
            await self._save_cover_image(data)

            # Get info from files
            path_gen = BookPathGenerator(library.root_path)
            parent = path_gen.generate_path(data)
            data["path"] = str(parent)
            data["library_id"] = library.id
            file_metadata = []

            for file in files:
                stats = await aios.stat(file)
                file_size = stats.st_size
                content_type, _ = mimetypes.guess_type(file)
                file_hash = await calculate_koreader_hash(file)

                filename = path_gen.generate_filename(data, Path(file.name))

                file_metadata.append(
                    FileMetadata(
                        path=str(filename),
                        size=file_size,
                        hash=file_hash,
                        content_type=content_type,
                    )
                )

            data["files"] = file_metadata

            # Move files to appropriate directory
            if len(files) > 1:
                await move_dir_contents(consume_path / group, parent)
            else:
                if len(group.parts) > 1:
                    await move_file(consume_path / group / files[0].name, parent / files[0].name)
                else:
                    await move_file(consume_path / group, parent / group)

            cleanup_empty_parent_directories(consume_path / group, consume_path)

            books.append(await super().create(data))
            await self.repository.session.commit()

        return books

    async def delete_books(
        self,
        book_ids: list[int],
        library: Library,
        delete_files: bool = False,
        **kwargs,
    ) -> None:
        """
        Delete books and optionally delete their associated files.

        Removes book records from the database and optionally deletes files from the
        filesystem, including cover images and book files.

        Args:
            book_ids: List of book IDs to delete.
            library: The library containing the books.
            delete_files: If True, also delete associated files from the filesystem.
            **kwargs: Additional keyword arguments.
        """
        books = await self.list(
            CollectionFilter(field_name="id", values=book_ids),
            Book.library_id == library.id,
        )

        for book in books:
            await self.remove_files(
                book.id,
                [file.id for file in book.files],
                library=library,
                delete_files=delete_files,
            )

            if book.cover_image:
                await delete_file(book.cover_image)

        await super().delete_many(book_ids)

    async def get_file(self, book_id: int, file_id: int) -> File:
        """
        Retrieve a file for download.

        Args:
            book_id: The ID of the book containing the file.
            file_id: The ID of the file to retrieve.

        Returns:
            A File object ready for download.

        Raises:
            ValueError: If the file is missing or not found for the given book.
        """
        book = await self.get(book_id)
        if book.path is None:
            raise ValueError("Cannot download file: book has no path")

        for file in book.files:
            if file.id != file_id:
                continue

            path = Path(book.path) / Path(file.path)
            if not await aios.path.isfile(path):
                raise ValueError("The file is missing")

            return File(path, media_type=file.content_type)

        raise ValueError("No such file for the given book")

    async def get_files(
        self, book_ids: list[int], library_id: int
    ) -> AsyncIterator[bytes]:
        """
        Get all selected book files as a compressed zip file.

        Streams the zip file in chunks to avoid loading the entire file into memory.

        Args:
            book_ids: List of book IDs to include in the zip.
            library_id: The ID of the library containing the books.

        Yields:
            Chunks of compressed zip file data.
        """
        books = await self.list(Book.id.in_(book_ids), Book.library_id == library_id)

        files = [file for book in books for file in book.files]

        buffer = BytesIO()
        with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
            for file in files:
                path = Path(file.path)
                if path.exists():
                    zip_file.write(path, arcname=path.name)

        buffer.seek(0)
        chunk_size = 32768  # 32 KiB
        while True:
            chunk = buffer.read(chunk_size)
            if not chunk:
                break
            yield chunk

    async def update_book(
        self, book_id: int, update_data: ModelDictT[Book], library: Library
    ) -> Book:
        """
        Update a book's metadata and files.

        Handles cover image updates, path reorganization based on updated metadata,
        and cleanup of empty directories. Ensures related entities (authors, series, etc.)
        remain unique in the database.

        Args:
            book_id: The ID of the book to update.
            update_data: The updated book data.
            library: The library containing the book.

        Returns:
            The updated Book entity.
        """
        book = await self.get(book_id)

        if not is_dict(update_data):
            data = service.schema_dump(update_data, exclude_unset=True)
        else:
            data = update_data

        data["id"] = book_id

        if data.get("cover_image", None):
            # Delete the existing cover image
            if book.cover_image:
                await delete_file(book.cover_image)
            await self._save_cover_image(data)

        # TODO: extract out into its own function _update_book_path
        # Check if file path must be updated (only for books with files)
        if book.path is not None:
            path_gen = BookPathGenerator(library.root_path)
            updated_path = path_gen.generate_path(book.to_dict() | data)
            if str(updated_path) != book.path:
                # TODO: Move only the files associated with the book instead of the whole directory
                await move_dir_contents(book.path, updated_path)
                data["path"] = str(updated_path)
                cleanup_empty_parent_directories(Path(book.path), Path(library.root_path))

        return await super().update(data, item_id=book_id, execution_options={"populate_existing": True})

    async def add_files(
        self, book_id: int, files: list[UploadFile], library: Library
    ) -> None:
        """
        Add additional files to an existing book.

        Args:
            book_id: The ID of the book.
            files: List of files to add.
            library: The library containing the book.
        """
        book = await self.get(book_id)
        data = book.to_dict()
        data["files"] = files
        new_files = await self._save_book_files(library, data)
        book.files.extend(new_files)
        await self.update_book(book.id, {"files": [file for file in book.files]}, library)

    async def remove_files(
        self, book_id: int, file_ids: list[int], delete_files: bool, library: Library
    ) -> None:
        """
        Remove files from an existing book.

        Optionally deletes the files from the filesystem and cleans up empty directories.

        Args:
            book_id: The ID of the book.
            file_ids: List of file IDs to remove.
            delete_files: If True, also delete the files from the filesystem.
            library: The library containing the book.
        """
        book = await self.get_one(Book.id == book_id, Book.library_id == library.id)

        if delete_files and book.path is not None:
            # TODO: Extract this out into its own function
            for file in (file for file in book.files if file.id in file_ids):
                full_path = Path(book.path) / Path(file.path)
                if await aios.path.isfile(full_path):
                    await aios.remove(full_path)

                cleanup_empty_parent_directories(
                    Path(book.path), Path(library.root_path)
                )

        book.files = [file for file in book.files if file.id not in file_ids]

        await super().update(book.files, book.id)

    async def to_model_on_create(self, data: ModelDictT[Book]) -> ModelDictT[Book]:
        """
        Hook called during model creation to populate unique relationships.

        Args:
            data: The book data to transform.

        Returns:
            The transformed data with unique relationship entities.
        """
        data = service.schema_dump(data)
        self._preprocess_book_data(data)
        return await self._populate_with_unique_relationships(data)

    async def to_model_on_update(self, data: ModelDictT[Book]) -> ModelDictT[Book]:
        """
        Hook called during model update to populate unique relationships.

        Args:
            data: The book data to transform.

        Returns:
            The transformed data with unique relationship entities.
        """
        data = service.schema_dump(data)
        self._preprocess_book_data(data)
        model_data = await self._populate_with_unique_relationships(data)

        return model_data
    
    def _preprocess_book_data(self, data: dict) -> dict:
        """Transform API input format to model format."""
        if not isinstance(data, dict):
            return data
        
        # Transform dict identifiers to list of Identifier objects
        if "identifiers" in data and isinstance(data["identifiers"], dict):
            data["identifiers"] = [
                Identifier(name=key, value=val)
                for key, val in data["identifiers"].items()
            ]
    
        return data

    async def _populate_with_unique_relationships(self, data: ModelDictT[Book]) -> ModelDictT[Book]:
        """
        Ensure relationship entities (authors, series, tags, etc.) are unique in the database.

        Fetches or creates unique instances of related entities to prevent duplicates.

        Args:
            data: Book data containing relationship fields.

        Returns:
            The data with relationship entities replaced by unique instances.
        """
        if not isinstance(data, dict):
            return data

        if data.get("id", None):
            model_data = await super().get(data["id"])

            # Update all scalar fields dynamically
            mapper = inspect(Book)
            for column in mapper.columns:
                if column.name in data:
                    setattr(model_data, column.name, data[column.name])
        else:
            model_data = await super().to_model(data)

        if "authors" in data:
            model_data.authors = [
                await Author.as_unique_async(self.repository.session, name=author)
                for author in data["authors"]
            ]

        if "series" in data:
            if data["series"]:
                model_data.series = await BookSeries.as_unique_async(
                    self.repository.session, title=data["series"]
                )
            else:
                model_data.series = None

        if "publisher" in data:
            if data["publisher"]:
                model_data.publisher = await Publisher.as_unique_async(
                    self.repository.session, name=data["publisher"]
                )
            else:
                model_data.publisher = None

        if "tags" in data:
            model_data.tags = [
                await Tag.as_unique_async(self.repository.session, name=tag)
                for tag in data["tags"]
            ]

        if "identifiers" in data:
            model_data.identifiers = data["identifiers"]

        if "files" in data:
            model_data.files = data["files"]

        return model_data

    async def _save_book_files(self, library: Library, data: dict) -> list[FileMetadata]:
        """
        Save uploaded book files to the filesystem.

        Uses the library's path template to organize files and extracts file metadata
        (size, content type, etc.) for database storage.

        Args:
            library: The library containing the book.
            data: Book data with files to save.

        Returns:
            The data with file paths and metadata populated.
        """
        # Use the library template path with the book's data to generate a filepath
        path_gen = BookPathGenerator(library.root_path)
        parent = path_gen.generate_path(data)
        data["path"] = str(parent)

        file_metadata = []
        CHUNK_SIZE = 262144  # 256 KiB
        for file in data.pop("files", []):
            if TYPE_CHECKING:
                file: UploadFile

            # TODO filename should be generated with generate_path function
            filename = path_gen.generate_filename(data, Path(file.filename).name)

            # Store files in the correct directory and return the file's metadata
            await file.seek(0)
            path = parent / filename
            path.parent.mkdir(parents=True, exist_ok=True)

            hasher = StreamingHasher()
            async with aiofiles.open(path, "wb") as dest:
                # Read spooled file and save it to the local filesystem
                while chunk := await file.read(CHUNK_SIZE):
                    await dest.write(chunk)
                    hasher.update(chunk)

            stats = await aios.stat(path)
            file_size = stats.st_size

            file_metadata.append(
                FileMetadata(
                    path=str(filename),
                    size=file_size,
                    hash=hasher.hexdigest(),
                    content_type=file.content_type,
                )
            )

        data["files"] = file_metadata
        return data["files"]

    async def _parse_metadata_from_files(self, data: dict, root_path: Path | None = None) -> dict:
        """
        Extract metadata (title, author, etc.) from book files.

        Fills in missing book metadata fields using extracted information from files.

        Args:
            data: Book data to populate with extracted metadata.

        Returns:
            The data with extracted metadata populated in empty fields.
        """
        extracted_metadata = await MetadataExtractor.extract_metadata(data["files"], root_path)

        # Add missing fields and update empty (falsey) fields with extracted metadata
        for attr in extracted_metadata.keys():
            if not data.get(attr, None):
                data[attr] = extracted_metadata[attr]
        return data

    async def _save_cover_image(self, data: dict) -> dict:
        """
        Save the book's cover image to the filesystem.

        Converts uploaded images to WebP format for consistent storage and assigns
        a unique UUID-based filename.

        Args:
            data: Book data potentially containing a cover_image field.

        Returns:
            The data with cover_image path populated or removed if not provided.
        """
        if (image := data.pop("cover_image", None)) is None:
            return data

        if isinstance(image, UploadFile):
            content = await image.read()
            image = Image.open(BytesIO(content))

        filepath = Path(settings.book_cover_path)
        filename = Path(f"{uuid.uuid4()}.webp")  # Random filename

        await save_image(image, filepath / filename)

        data["cover_image"] = str(filepath / filename)
        return data
