import enum
import logging
from functools import lru_cache
from hashlib import sha256
from pathlib import Path
from typing import Literal, assert_never, get_args

from pydantic import (
    BaseModel,
    Field,
)

from mcpunk.util import matches_filter

logger = logging.getLogger(__name__)


class ChunkCategory(enum.StrEnum):
    callable = "callable"
    markdown_section = "markdown section"
    imports = "imports"
    module_level = "module_level"
    whole_file = "whole_file"
    other = "other"


# Seems if you annotate a FastMCP tool function with an enum it totally
# crashes claude desktop. So define an equivalent Literal type here.
ChunkCategoryLiteral = Literal[
    "callable",
    "markdown section",
    "imports",
    "module_level",
    "whole_file",
    "other",
]
assert set(get_args(ChunkCategoryLiteral)) == set(ChunkCategory.__members__.values())


class Chunk(BaseModel):
    """A chunk of a file, e.g. a function or a markdown section."""

    category: ChunkCategory = Field(description="`function`, `markdown section`, `imports`")
    name: str = Field(description="`my_function` or `MyClass` or `# My Section`")
    line: int | None = Field(description="Line within file where it starts. First line is 1.")
    content: str = Field(description="Content of the chunk")

    def id_(self, path: Path | None) -> str:
        """Generate a (probably) unique ID based on content, name, line, and category.

        This approach means that the ID will stay the same even if the chunk is recreated
        (e.g. as opposed to a totally random ID).
        The id includes the chunk name, which makes debugging and monitoring of tool
        requests far nicer.

        Incorporates an optional path argument to make the id more unique, especially
        across files.

        This does some shenanigans with caching and pre-hashing things
        to help makes things faster.

        Risk of collisions and so on, unlikely, but probably not astronomically unlikely :D
        """
        # Use Python's built-in hash which is much faster than SHA256. Then clip
        # that to a 64-bit int.
        content_hash = str(hash(self.content) % 0xFFFFFFFFFFFFFFFF)
        return _cached_hash(
            content_hash=content_hash,
            name=self.name,
            line=self.line,
            category=self.category,
            path=path,
        )

    def matches_filter(
        self,
        filter_: None | list[str] | str,
        filter_on: Literal["name", "content", "name_or_content"],
    ) -> bool:
        """Return True if the chunk's name matches the given filter.

        str matches if the chunk's name contains the string.
        list[str] matches if the chunk's name contains any of the strings in the list.
        None matches all chunks.
        """
        if filter_on == "name":
            data = self.name
        elif filter_on == "content":
            data = self.content
        elif filter_on == "name_or_content":
            data = self.content + self.name
        else:
            assert_never(filter_on)
        return matches_filter(filter_, data)

    def split(
        self,
        max_size: int = 10_000,
        split_chunk_prefix: str = (
            "[This is a subsection of the chunk. Other parts contain the rest of the chunk]\n\n"
        ),
    ) -> list["Chunk"]:
        """Split this chunk into smaller chunks.

        This will split the chunk at line boundaries, unless the
        line is already longer than max_size.

        Args:
            max_size: Maximum size in characters for the chunk contents. At least 100.
            split_chunk_prefix: Prefix to add the start of each newly created split chunk.
                Unused if the chunk is not split. You can set to empty string to
                suppress the prefix.

        Returns:
            List containing either the original chunk (if small enough) or multiple smaller chunks
        """
        assert max_size >= 100, "max_size must be at least 100"
        # If chunk is small enough, return it as is
        if len(self.content) <= max_size:
            return [self]
        max_size -= len(split_chunk_prefix)
        assert max_size > 0, f"{max_size} maybe decrease prefix length"

        result: list[Chunk] = []
        max_line_size = max_size - 50  # Leave some margin

        # Preprocess to split long lines first. This could be avoided, but it does
        # make the whole thing a bit simpler as we always know later on that a single line
        # will never be longer than max_size.
        processed_lines = []
        for line in self.content.splitlines(keepends=True):
            if len(line) > max_line_size:
                # Split the line into chunks of max_line_size
                for i in range(0, len(line), max_line_size):
                    processed_lines.append(line[i : i + max_line_size])
            else:
                processed_lines.append(line)

        # Now split into chunks of max_size
        current_content: list[str] = []
        current_size = 0
        part_num = 1

        for line in processed_lines:
            # If adding this line would exceed the limit, create a new chunk
            if current_size + len(line) > max_size and current_content:
                new_chunk = Chunk(
                    category=self.category,
                    name=f"{self.name}_part{part_num}",
                    content=split_chunk_prefix + "".join(current_content),
                    line=None,
                )
                result.append(new_chunk)
                part_num += 1
                current_content = []
                current_size = 0

            # Add the line to the current chunk
            current_content.append(line)
            current_size += len(line)

        # Add the final chunk if there's anything left
        if current_content:
            new_chunk = Chunk(
                category=self.category,
                name=f"{self.name}_part{part_num}",
                content=split_chunk_prefix + "".join(current_content),
                line=None,
            )
            result.append(new_chunk)

        return result


@lru_cache(maxsize=100_000)
def _cached_hash(
    *,
    content_hash: str,
    name: str,
    line: int | None,
    category: ChunkCategory,
    path: Path | None,
) -> str:
    components = [
        content_hash,
        name,
        str(line),
        str(category),
        str(path),
    ]
    return name + "_" + sha256("".join(components).encode()).hexdigest()[:10]
