# src/chitai/services/metadata_extractor.py

# TODO: Code is a mess. Clean it up and add docstrings

# Standard library
from abc import ABC, abstractmethod
import datetime
from pathlib import Path
from io import BytesIO
import re
from tempfile import SpooledTemporaryFile
from typing import Any, Protocol, BinaryIO
import logging

# Third-party libraries
import PIL
import PIL.Image
from litestar.datastructures import UploadFile
import pypdfium2
import ebooklib
from ebooklib import epub

# Local imports
from chitai.services.utils import (
    create_temp_file,
    get_file_extension,
    get_filename,
    is_valid_isbn,
)

logger = logging.getLogger(__name__)


class FileExtractor(Protocol):
    @classmethod
    async def extract_metadata(
        cls, input: UploadFile | BinaryIO | bytes | Path | str
    ) -> dict[str, Any]: ...

    @classmethod
    async def extract_text(cls, input: UploadFile | BinaryIO | bytes | Path | str) -> str: ...


class Extractor:
    """Default extractor class that delegates to the proper extractor based on file type."""

    format_priorities = {"epub": 1, "pdf": 2}

    @classmethod
    async def extract_metadata(cls, files: list[UploadFile] | list[Path], root_path: Path | None = None) -> dict[str, Any]:
        metadata = {}

        # Sort based on file priority
        # EPUB tends to give better metadata results over pdf
        sorted_files = sorted(files, key=lambda f: Extractor._get_file_priority(f))
        
        for file in sorted_files:
            match get_file_extension(file):
                case "epub":
                    metadata = metadata | await EpubExtractor.extract_metadata(file)
                case "pdf":
                    metadata = metadata | await PdfExtractor.extract_metadata(file)
                case _:
                    break

        # Get metadata from file names
        for file in files:
            metadata = FilenameExtractor.extract_metadata(file) | metadata

        # Get metadata from filepath
        metadata = metadata | FilepathExtractor.extract_metadata(files[0], root_path)

        # format the title
        if metadata.get('title', None):
            title, subtitle = Extractor.format_book_title(metadata["title"])
            metadata["title"] = title
            metadata["subtitle"] = subtitle

        return metadata

    @classmethod
    async def extract_text(cls, input: BinaryIO | bytes | Path | str) -> str: ...

    @classmethod
    def format_book_title(cls, title: str) -> tuple[str, str | None]:
        # Convert all trailing underscores to colons
        title = Extractor.convert_underscores_to_colons(title)

        colon_count = title.count(":")

        if colon_count < 2:
            return title, None  # No subtitle

        # There is a subtitle, split the string at the second colon (colon is disregarded)
        first_colon_idx = title.find(":")
        second_colon_idx = title.find(
            ":", first_colon_idx + 1
        )  # Start search after first colon

        return (title[:second_colon_idx], title[second_colon_idx + 1 :])

    @classmethod
    def _get_file_priority(cls, upload: UploadFile | Path) -> float:

        filename = ""
        if isinstance(upload, UploadFile):
            filename = Path(upload.filename)
        elif isinstance(upload, Path):
            filename = upload.name

        file_ext = get_file_extension(filename)

        if file_ext is None:
            return float('inf')
        
        return Extractor.format_priorities.get(file_ext, float("inf"))

    @classmethod
    def convert_underscores_to_colons(cls, text: str) -> str:
        """
        Converts underscores that trail a word and are followed by a space and another word to colons.
        Only converts when there's a space between words.

        For example:
        - "hello_ world" -> "hello: world"
        - "this_ is_ a_ test" -> "this: is: a: test"
        - "hello_world" -> "hello_world" (unchanged, no space)
        - "word_" -> "word_" (unchanged, no following word)

        Args:
            text (str): The input text to be processed

        Returns:
            str: The processed text with underscores converted to colons
        """
        # This pattern matches a word followed by underscore(s), followed by a space, followed by another word
        pattern = r"(\w+)(_+)(\s+)(\w+)"

        # Replace all occurrences with word1: word2 (preserving the space)
        result = re.sub(pattern, r"\1:\3\4", text)

        return result


class PdfExtractor(FileExtractor):
    """Extract metadata and text from a PDF file."""

    @classmethod
    async def extract_metadata(
        cls, data: UploadFile | SpooledTemporaryFile | BinaryIO | bytes | Path | str
    ) -> dict[str, Any]:
        """Extracts metadata from a PDF file."""
        metadata = {}

        if isinstance(data, UploadFile):
            data = data.file

        doc = pypdfium2.PdfDocument(data)

        basic_metadata = doc.get_metadata_dict(skip_empty=False)
        metadata["title"] = basic_metadata["Title"]
        metadata["authors"] = PdfExtractor._parse_authors(basic_metadata["Author"])
        metadata["pages"] = len(doc)
        # Use the documents creation date as a guess for the date published
        metadata["published_date"] = PdfExtractor._parse_pdf_date(
            basic_metadata["CreationDate"]
        )
        metadata["cover_image"] = PdfExtractor._extract_cover(doc)

        text_by_page = PdfExtractor._extract_text_from_pages(doc, 6)
        metadata["identifiers"] = PdfExtractor._extract_isbns("".join(text_by_page))

        doc.close()

        return {key: val for key, val in metadata.items() if val}  # Remove empty values

    @classmethod
    async def extract_text(cls, data: BinaryIO | bytes | Path | str) -> str:
        raise NotImplementedError()  # TODO: Implement text extraction from PDF files

    @classmethod
    def _parse_authors(cls, author_string: str):
        """Extracts author data from the authors string based on common formats."""

        # Remove trailing delimiters and whitespace
        author_string = author_string.strip("; &").strip()

        if not author_string:
            return []

        try:
            # Determine the format based on presence of commas
            if "," in author_string:
                # Handle LastName, FirstName; LastName, FirstName format
                authors = [a.strip() for a in author_string.split(";") if a.strip()]
                formatted_authors = []
                for author in authors:
                    last_name, first_name = [name.strip() for name in author.split(",")]
                    formatted_authors.append(f"{first_name} {last_name}")
            elif "&" in author_string:
                # Handle FirstName LastName & FirstName LastName format
                authors = [a.strip() for a in author_string.split("&") if a.strip()]
                formatted_authors = authors

            elif ";" in author_string:
                # Handle FirstName LastName; FirstName LastName format
                authors = [a.strip() for a in author_string.split(";") if a.strip()]
                formatted_authors = authors

            else:
                return [author_string]

            return formatted_authors

        except Exception:
            return []

    @classmethod
    def _parse_pdf_date(cls, date_string: str | None) -> datetime.date | None:
        if date_string is None:
            return None

        # Remove "D:" prefix if present
        if date_string.startswith("D:"):
            date_string = date_string[2:]

        # Extract just the date portion (YYYYMMDD) and parse it
        date_portion = date_string[:8]  # Just take first 8 characters (YYYYMMDD)
        try:
            return datetime.datetime.strptime(date_portion, "%Y%m%d").date()

        except Exception as e:
            return None

    @classmethod
    def _extract_text_from_pages(cls, doc: pypdfium2.PdfDocument, num_pages=5):
        num_pages = min(num_pages, len(doc))

        text_by_page = []

        for i in range(num_pages):
            text_by_page.append(doc[i].get_textpage().get_text_bounded())

        return text_by_page

    @classmethod
    def _extract_isbns(cls, text: str) -> dict[str, str]:
        isbn_pattern = (
            r"(\b\d{9}[\dXx]\b|\b\d{1,5}-?\d{1,7}-?\d{1,6}-?\d{1,6}-?[\dXx]\b)"
        )

        isbns = {}

        matches = re.findall(isbn_pattern, text)
        for match in matches:
            # Strip hyphens if present
            isbn = match.replace("-", "")

            if not is_valid_isbn(isbn):
                continue

            if len(isbn) == 10:
                isbns.update({"isbn-10": isbn})
            elif len(isbn) == 13:
                isbns.update({"isbn-13": isbn})

            if len(isbns) >= 2:  # Exit after both isbns have been matched
                break

        return isbns

    @classmethod
    def _extract_cover(cls, doc: pypdfium2.PdfDocument) -> PIL.Image.Image | None:
        page = doc[0]
        bitmap: pypdfium2.PdfBitmap = page.render()
        image = bitmap.to_pil()
        return image


class EpubExtractor(FileExtractor):
    """Extract metadata and text from an EPUB file."""

    @classmethod
    async def extract_metadata(
        cls, input: UploadFile | BinaryIO | bytes | Path | str
    ) -> dict[str, Any]:
        """Extract metadata from an EPUB file.

        Args:
            input: Can be an UploadFile, file-like object, bytes, Path, or string path

        Returns:
            Dictionary containing extracted metadata
        """

        # Handle different input types
        if isinstance(input, UploadFile):
            await input.seek(0)
            data = await input.read()
            path = await create_temp_file(data)
        elif isinstance(input, bytes):
            path = await create_temp_file(input)
        elif isinstance(input, BinaryIO):
            input.seek(0)
            data = input.read()
            path = await create_temp_file(data)
        elif isinstance(input, (Path, str)):
            path = Path(input) if isinstance(input, str) else input
        else:
            raise TypeError(f"Unsupported input type: {type(input)}")

        try:
            metadata = EpubExtractor._get_metadata(path)
        except Exception as e:
            logger.error(f"Error extracting metadata from epub: {e}")
            return {}
        return metadata

    @classmethod
    def extract_text(cls, input: BinaryIO | bytes | Path | str) -> str:
        raise NotImplementedError()  # TODO: Implement text extraction from EPUB files

    @classmethod
    def _get_metadata(cls, file_path) -> dict[str, Any]:
        """Uses ebooklib to extract metadata from an EPUB file.

        TODO: Some EPUBs fail to parse due to missing files, this is an issue with ebooklib;
                more info and possible solution here: https://github.com/aerkalov/ebooklib/issues/281
        """
        book = epub.read_epub(file_path)
        metadata = {}

        metadata["title"] = book.get_metadata("DC", "title")[0][0]
        metadata["identifiers"] = EpubExtractor._extract_identifiers(book)
        metadata["language"] = book.get_metadata("DC", "language")[0][0]
        metadata["published_date"] = EpubExtractor._extract_published_date(book)
        metadata["description"] = EpubExtractor._extract_description(book)
        metadata["publisher"] = EpubExtractor._extract_publisher(book)
        metadata["authors"] = EpubExtractor._extract_authors(book)
        metadata["cover_image"] = EpubExtractor._extract_cover(book)

        metadata = {
            key: val for key, val in metadata.items() if val
        }  # Remove all empty values

        return metadata

    @classmethod
    def _extract_authors(cls, epub: epub.EpubBook) -> list[str]:
        authors = []

        for creator in epub.get_metadata("DC", "creator"):
            authors.append(creator[0])

        return authors

    @classmethod
    def _extract_identifiers(cls, epub: epub.EpubBook) -> dict[str, str]:
        identifiers = {}

        for id in epub.get_metadata("DC", "identifier"):
            if is_valid_isbn(id[0]):
                if len(id[0]) == 13:
                    identifiers.update({"isbn-13": id[0]})

                elif len(id[0]) == 10:
                    identifiers.update({"isbn-10": id[0]})

        return identifiers

    @classmethod
    def _extract_description(cls, epub: epub.EpubBook) -> str | None:
        try:
            return epub.get_metadata("DC", "description")[0][0]

        except:
            return None

    @classmethod
    def _extract_published_date(cls, epub: epub.EpubBook) -> datetime.date | None:
        try:
            date_str = epub.get_metadata("DC", "date")[0][0].split("T")[0]
            return datetime.date.fromisoformat(date_str)

        except:
            return None

    @classmethod
    def _extract_publisher(cls, epub: epub.EpubBook) -> str | None:
        try:
            epub.get_metadata("DC", "publisher")[0][0]

        except:
            return None

    @classmethod
    def _extract_cover(cls, epub: epub.EpubBook) -> PIL.Image.Image | None:
        """Extract cover image from EPUB."""

        # Strategy 1: Check for cover metadata (most common)
        try:
            cover_meta = epub.get_metadata("OPF", "cover")
            if cover_meta:
                # Try to get the cover ID from metadata
                for meta in cover_meta:
                    if isinstance(meta, tuple) and len(meta) > 1:
                        # Handle both dict and direct string references
                        cover_id = (
                            meta[1].get("content")
                            if isinstance(meta[1], dict)
                            else meta[1]
                        )
                        if cover_id:
                            cover_item = epub.get_item_with_id(cover_id)
                            if cover_item:
                                return PIL.Image.open(BytesIO(cover_item.content))
        except Exception as e:
            pass  # Fallback to next strategy

        # Strategy 2: Search image filenames for "cover" keyword
        images = list(epub.get_items_of_type(ebooklib.ITEM_IMAGE))
        for image in images:
            filename = (
                str(image.get_name()).casefold()
                if hasattr(image, "get_name")
                else str(image).casefold()
            )
            if "cover" in filename:
                return PIL.Image.open(BytesIO(image.content))

        return None


class FilepathExtractor(FileExtractor):
    """Extracts metadata from the filepath."""

    @classmethod
    def extract_metadata(cls, input: UploadFile | Path | str, root_path: Path | None = None) -> dict[str, Any]:
        
        if isinstance(input, UploadFile):
            path = Path(input.filename).parent
        else:
            path = Path(input).parent

        if root_path:
            path = path.relative_to(root_path)
        
        parts = path.parts
        metadata: dict[str, str | None] = {}
        
        if len(parts) == 3:
            # Format: Author/Series/Part - Title/filename
            metadata['author'] = parts[0]
            
            # Extract part number and title from directory name (parts[2])
            dirname = parts[2]
            match = re.match(r'^([\d.]+)\s*-\s*(.+)$', dirname)
            
            if match:
                metadata['series_position'] = match.group(1)  # Keep as string
                metadata['series'] = parts[1]
                metadata['title'] = match.group(2).strip()
            else:
                metadata['series'] = parts[1]
                metadata['title'] = path.stem
                
        elif len(parts) == 2:
            # Format: Author/Title
            metadata['author'] = parts[0]
            metadata['title'] = path.stem  # Remove extension
        
        return metadata

class FilenameExtractor(FileExtractor):
    """Extracts metadata from the filename."""

    @classmethod
    def extract_metadata(cls, input: Path | str | UploadFile) -> dict[str, Any]:

        if isinstance(input, str):
            filename = get_filename(Path(input), ext=False)
        elif isinstance(input, Path):
            filename = get_filename(input, ext=False)
        elif isinstance(input, UploadFile):
            filename = Path(input.filename).name
        else:
            raise ValueError("Input type not supported")

        # Common patterns to split title and authors
        patterns = [
            r"^(.+?)\s*[-–]\s*([^-–]+)$",  # Pattern with hyphen separator
            r"^(.+?)\s*by\s+([^-–]+)$",  # Pattern with "by" separator
            r"^([^:]+):\s*(.+)$",  # Pattern with colon separator
        ]

        title = filename
        authors_str = ""

        # Try each pattern until we find a match
        for pattern in patterns:
            match = re.match(pattern, filename, re.IGNORECASE)
            if match:
                title, authors_str = match.groups()
                break

        # Clean up title
        title = re.sub(r"\s*\d+[^a-zA-Z]+.*$", "", title)  # Remove edition numbers
        title = re.sub(r"\s*\([^)]*\)", "", title)  # Remove parentheses
        title = title.strip()

        # Parse authors
        authors = []
        if authors_str:
            # Split on common author separators
            author_parts = re.split(r"[,_]|\sand\s|\s&\s", authors_str)

            for author in author_parts:
                # Clean up each author name
                author = author.strip()
                if author:
                    # Remove any remaining parenthetical content
                    author = re.sub(r"\s*\([^)]*\)", "", author)
                    # Remove titles and other common prefixes
                    author = re.sub(
                        r"^(MD\.Dr\.|Mr\.|Mrs\.|Ms\.|Prof\.)\s+", "", author
                    )
                    authors.append(author.strip())

        return {"title": title, "authors": authors}

    @classmethod
    def extract_text(cls, input: Path | str):
        raise NotImplementedError()  # Obviously can't extract text from filename...
