import asyncio
import base64
import contextlib
import difflib
import os
import tempfile
import xml.etree.ElementTree as ET
from typing import Annotated, ClassVar, Literal, TypedDict, cast, override
from pathlib import Path as StdPath
from string import Template
from anyio import Path as AnyioPath
from anyxml import AnyXml
from loguru import logger
from dais_sdk.types import AudioBlock, Base64Source, ContentBlock, ImageBlock, VideoBlock
from dais_scantree import bfs as scantree_bfs, dfs as scantree_dfs
from dais_scantree.ignore_rule import load_gitignore_spec
from wcmatch import glob as wc_glob
from src.db import db_context
from src.db.models import toolset as toolset_models
from src.services.markdown_cache import MarkdownCacheService
from src.binaries import RIPGREP_PATH

from src.utils import MarkdownConverter
from ..toolset_wrapper import builtin_tool, BuiltinToolset, BuiltinToolsetContext, BuiltinToolDefaults
from ...utils import magika_identify_path


MAX_MEDIA_CONTENT_BLOCK_BYTES = 50 * 1024 * 1024


# Since `is_binary` from binaryornot sometimes misdetects some files as binary,
# we defines a enhanced wrapper function here.
# TODO: remove this function once binaryornot.check.is_binary is fixed.
def is_binary(path: StdPath) -> bool:
    from binaryornot.helpers import has_binary_extension, is_binary_string  # pyright: ignore[reportAttributeAccessIssue]
    if has_binary_extension(path):
        return True
    with open(path, "rb") as f:
        bytes = f.read()
        return is_binary_string(bytes)

class PathExpander:
    def __init__(self, path_envs: dict[str, str]):
        self._path_envs = path_envs

    def substitute(self, path: str) -> str:
        template = Template(path)
        return template.safe_substitute(self._path_envs)

class FileSystemToolset(BuiltinToolset):
    _logger = logger.bind(name="FileSystemToolset")
    _file_locks: ClassVar[dict[str, asyncio.Lock]] = {}
    _file_locks_lock: ClassVar[asyncio.Lock] = asyncio.Lock()

    def __init__(self,
                 ctx: BuiltinToolsetContext,
                 toolset_ent: toolset_models.Toolset | None = None):
        from ...notes import NoteMaterializer
        from ...skills import SkillMaterializer

        super().__init__(ctx, toolset_ent)
        self._path_expander = PathExpander({
            **SkillMaterializer.get_skill_dir_env(),
            **NoteMaterializer.get_notes_dir_env(ctx.workspace_id),
        })
        self._markdown_converter = MarkdownConverter()

    @property
    @override
    def name(self) -> str: return "FileSystem"

    def _resolve_path(self, file_path: str) -> AnyioPath:
        expanded_path = self._path_expander.substitute(file_path)
        expanded_path = os.path.expanduser(expanded_path)
        return AnyioPath(self._ctx.cwd / expanded_path)

    async def _get_file_lock(self, path: StdPath | AnyioPath) -> asyncio.Lock:
        lock_key_path = await AnyioPath(path).resolve(strict=False)
        normalized_key = str(lock_key_path)
        async with self._file_locks_lock:
            lock = self._file_locks.get(normalized_key)
            if lock is None:
                lock = asyncio.Lock()
                self._file_locks[normalized_key] = lock
            return lock

    async def _atomic_write_text(self, path: AnyioPath, content: str) -> None:
        def write_impl() -> None:
            std_path = StdPath(path)
            if std_path.is_symlink():
                std_path = std_path.resolve(strict=False)

            temp_file_path: str | None = None
            try:
                with tempfile.NamedTemporaryFile(
                    "w",
                    encoding="utf-8",
                    dir=std_path.parent,
                    delete=False,
                ) as temp_file:
                    temp_file_path = temp_file.name
                    temp_file.write(content)
                    temp_file.flush()
                    os.fsync(temp_file.fileno())
                os.replace(temp_file_path, std_path)
            finally:
                if temp_file_path is not None:
                    with contextlib.suppress(FileNotFoundError):
                        os.unlink(temp_file_path)

        await asyncio.to_thread(write_impl)

    @builtin_tool(validate=True, defaults=BuiltinToolDefaults(auto_approve=True))
    async def read_file(self,
                        path: Annotated[str,
                          "The path of the file to read (relative to the current working directory)."],
                        offset: Annotated[int, "The line number to start reading from (1-based)."] = 1,
                        max_lines: Annotated[int, "The maximum number of lines to read."] = 2000,
                        ) -> str | list[ContentBlock]:
        """
        Read the contents of a file at the specified path.
        For image, audio, and video files, this tool will return the file as a ContentBlock array;
        for text files, this tool will directly return the file content;
        for .pdf, .docx, .pptx, .xlsx, .epub files, this tool will convert the file to markdown format and return the markdown text.

        Returns:
            A XML string with the file content and metadata as attributes:
            - start_line: The line number of the first line returned (1-based)
            - end_line: The line number of the last line returned (1-based)
            - total_lines: The total number of lines in the file

            Example:
                <file_content start_line="1" end_line="50" total_lines="312">
                def foo():
                    ...
                </file_content>

        To read the next chunk, pass end_line + 1 as the offset in the next call.
        """
        async def convert_to_markdown_with_cache(path: AnyioPath) -> str:
            async with db_context() as db_session:
                markdown_cache_service = MarkdownCacheService(db_session, self._ctx.workspace_id, self._ctx.cwd)
                cached = await markdown_cache_service.get(path)
                if cached is not None: return cached

                converted = await self._markdown_converter.convert(path)
                await markdown_cache_service.set(path, converted)
                return converted

        def format_result(content: str,
                          total_lines: int,
                          start_line: int | None = None,
                          end_line: int | None = None) -> str:
            attributes = {"total_lines": str(total_lines)}
            if start_line is not None:
                attributes["start_line"] = str(start_line)
            if end_line is not None:
                attributes["end_line"] = str(end_line)

            root = ET.Element("file_content", attrib=attributes)
            root.text = AnyXml.RawText(content)
            return AnyXml.tostring(root)

        async def read_media_content_blocks(path: AnyioPath,
                                            media_type: Literal["image", "audio", "video"],
                                            mime_type: str) -> list[ContentBlock]:
            file_size = (await path.stat()).st_size
            if file_size > MAX_MEDIA_CONTENT_BLOCK_BYTES:
                raise ValueError(
                    f"File {path.name} is too large to return as a ContentBlock "
                    f"({file_size} bytes, max {MAX_MEDIA_CONTENT_BLOCK_BYTES} bytes)."
                )

            file_bytes = await path.read_bytes()
            encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
            source = Base64Source(mime_type=mime_type, data=encoded.decode("ascii"))
            match media_type:
                case "image": return [ImageBlock(source=source)]
                case "audio": return [AudioBlock(source=source)]
                case "video": return [VideoBlock(source=source)]

        abs_path = self._resolve_path(path)

        if not await abs_path.exists():
            raise FileNotFoundError(f"File not found at {path}")

        if self._markdown_converter.is_convertable_binary(abs_path):
            result = await convert_to_markdown_with_cache(abs_path)
            lines = result.splitlines()
        elif (output := await magika_identify_path(abs_path)) and output.group in {"image", "audio", "video"}:
            media_type, mime_type = output.group, output.mime_type
            return await read_media_content_blocks(abs_path, cast(Literal["image", "audio", "video"], media_type), mime_type)
        elif is_binary(StdPath(abs_path)):
            raise ValueError(f"File {path} is a binary file, and is not supported to read.")
        else:
            lines = (await abs_path.read_text("utf-8")).splitlines()

        result_lines = lines[(offset - 1):(offset + max_lines - 1)]
        return format_result(
            "\n".join(result_lines),
            start_line=offset,
            end_line=offset + len(result_lines) - 1,
            total_lines=len(lines))

    @builtin_tool(validate=True)
    async def write_file(self,
                         path: Annotated[str,
                          "The path of the file to write (relative to the current working directory)."],
                         content: Annotated[str,
                          "The content to write to the file."],
                         append: Annotated[bool,
                          "If True, content will be appended to the end of the file instead of overwriting it."] = False,
                         ) -> str:
        """
        Write content to a file at the specified path.
        Use this when you need to create a new file or overwrite an existing file with new content.
        If the parent directory of the specified path does not exist, it will be created automatically.

        IMPORTANT: This tool overwrites existing files. Before using it, ensure the target file does not contain important information that should be preserved.

        Examples:
            >>> write_file("test.txt", "Hello World!")
            File written successfully.

        Returns:
            A success message if the file was written successfully.
        """
        abs_path = self._resolve_path(path)
        lock = await self._get_file_lock(abs_path)
        async with lock:
            await abs_path.parent.mkdir(parents=True, exist_ok=True)

            if append:
                async with await abs_path.open("a", encoding="utf-8") as f:
                    await f.write(content)
            else:
                await self._atomic_write_text(abs_path, content)
        return "File written successfully."

    @builtin_tool(validate=True)
    async def edit_file(self,
                        path: Annotated[str,
                            "The path of the file to edit (relative to the current working directory)."],
                        old_content: Annotated[str,
                            "The exact snippet to be replaced. Must match exactly once in the file."],
                        new_content: Annotated[str,
                            "The new snippet that will replace old_content."],
                        expected_replacements: Annotated[int,
                            "Number of replacements expected. Use when you want to replace multiple occurrences of the same text."] = 1
                        ) -> str:
        """
        Edit a file by replacing a specific snippet with new content.
        `old_content` and `new_content` should be the MINIMAL snippet necessary to make the change (though you may include a few extra lines BEFORE and AFTER the target text to ensure uniqueness).

        Use this when you need to edit an existing file with existing content you have read before.
        If the target file is a binary file, the tool will raise an error without editing it.
        If you intend to replace multiple identical occurrences at once, set `expected_replacements` to the exact count.
        The tool will raise an error if the actual count is fewer than `expected_replacements`.

        Returns:
            The unified diff of the applied change.

        Note:
            If old_content matches more times than expected, expand it to include more surrounding context until it uniquely identifies the target location(s).
        """
        def generate_diff(old_content: str, new_content: str, file_name: str) -> str:
            diff = difflib.unified_diff(
                old_content.splitlines(),
                new_content.splitlines(),
                fromfile=f"a/{file_name}",
                tofile=f"b/{file_name}",
            )
            return "\n".join(diff)

        abs_path = self._resolve_path(path)
        lock = await self._get_file_lock(abs_path)
        async with lock:
            if not await abs_path.exists():
                raise FileNotFoundError(f"File not found at {path}")
            if is_binary(StdPath(abs_path)):
                raise ValueError(f"File {path} is a binary file, and is not supported to edit.")

            try:
                content = await abs_path.read_text("utf-8")
            except UnicodeDecodeError as exc:
                raise ValueError(f"File {path} is a binary file, and is not supported to edit.") from exc
            count = content.count(old_content)
            if count == 0:
                raise ValueError(f"Content not found in file: {path}")
            if count < expected_replacements:
                raise ValueError(
                    f"Expected {expected_replacements} occurrence(s) of the given content in {path}, "
                    f"but found {count}. Adjust `old_content` or `expected_replacements` accordingly."
                )

            old_file_content = content
            new_file_content = content.replace(old_content, new_content, expected_replacements)
            diff = await asyncio.to_thread(generate_diff, old_file_content, new_file_content, path)
            await self._atomic_write_text(abs_path, new_file_content)
        return diff

    def _list_directory_impl(self,
                             path: str,
                             max_depth: int | None,
                             show_all: bool,
                             ) -> str:
        """
        Since the scan logic will be slower after refactored to async version,
        we leave it sync here.
        """
        def format_item(item: StdPath) -> str:
            if item.is_file(follow_symlinks=False):
                return item.name
            elif item.is_dir(follow_symlinks=False):
                return item.name + "/"
            elif item.is_symlink():
                try:
                    target = item.readlink()
                except (OSError, ValueError):
                    return f"{item.name} -> <unreadable>"
                target_path = target.as_posix()
                if (item.parent / target).is_dir():
                    target_path += "/"
                return f"{item.name} -> {target_path}"
            else:
                return "<special file>"

        def format_items_flat(directory: StdPath) -> list[str]:
            """Format directory contents in non-recursive mode."""
            gitignore_spec = (load_gitignore_spec(abs_path)
                             if not include_gitignored else None)

            items: list[StdPath] = []
            for entry in directory.iterdir():
                if not include_hidden and entry.name.startswith("."):
                    continue
                entry_path = entry.name
                if entry.is_dir():
                    entry_path += "/"
                if gitignore_spec is not None:
                    if gitignore_spec.match_file(entry_path):
                        continue
                items.append(entry)

            items = sorted(items, key=lambda x: (not x.is_dir(), x.name.lower()))
            if len(items) == 0:
                return ["(empty directory)"]

            lines = []
            for idx, item in enumerate(items, 1):
                formated = format_item(item)
                lines.append(f"{idx}. {formated}")
            return lines

        def format_items_recursive(directory: StdPath, max_depth: int | None) -> list[str]:
            """
            Recursively format directory contents.

            Args:
                directory: Current directory path
                max_depth: Maximum depth limit
            """
            lines = []
            for entry, depth in scantree_dfs(directory, MAX_SCAN_LIMIT, max_depth, include_hidden, include_gitignored):
                indent = "  " * (depth - 1)
                entry_path = StdPath(entry)
                lines.append(f"{indent}{format_item(entry_path)}")
            return lines

        if max_depth is not None and max_depth < 1:
            raise ValueError(f"Invalid max_depth: {max_depth}")

        MAX_SCAN_LIMIT = 200_000
        abs_path = StdPath(self._resolve_path(path))
        include_hidden = show_all
        include_gitignored = show_all

        if not abs_path.exists():
            raise FileNotFoundError(f"Directory not found at {path}")
        if not abs_path.is_dir():
            raise NotADirectoryError(f"Path {path} is not a directory")

        result_lines = [f"Directory: {abs_path}"]
        if max_depth == 1:
            result_lines.extend(format_items_flat(abs_path))
        else:
            result_lines.extend(format_items_recursive(abs_path, max_depth))
        return "\n".join(result_lines)

    @builtin_tool(validate=True, defaults=BuiltinToolDefaults(auto_approve=True))
    async def list_directory(self,
                             path: Annotated[str,
                              "The path of the directory to list contents for (relative to the current working directory)."] = ".",
                             max_depth: Annotated[int | None,
                              """
                              Maximum depth for directory listing.
                              - 1 (default): Only direct children, listed flat with numbering
                              - 2: List up to 2 levels deep, with indentation
                              - n: List up to n levels deep, with indentation
                              - None: No depth limit, list all nested directories
                              """] = 1,
                             show_all: Annotated[bool,
                              "Whether to include hidden files and files ignored by .gitignore. "
                              "Use this if you can't find a specific file you're looking for."] = False,
                             ) -> str:
        """
        List files and directories within the specified directory.

        Use this when you need to explore the project structure, understand the codebase organization, or find specific files.
        Directories are indicated by a trailing slash. Symlinks show their target path.

        Examples:
            # Non-recursive listing:
            >>> list_directory("src")
            Directory: /absolute/path/to/src
            1. agent/
            2. db/
            3. external_lib -> /usr/local/lib/
            4. __init__.py
            5. app.py

            - - -

            # Recursive listing with depth limit:
            >>> list_directory("src", max_depth=2)
            Directory: /absolute/path/to/src
            routes/
              endpoint1/
              endpoint2/
            schemas/
            main.py
            broken_link -> <unreadable>
            __init__.py

            - - -

            # Unlimited recursive listing:
            >>> list_directory("src", max_depth=None)
            Directory: /absolute/path/to/src
            routes/
              endpoint1/
                subendpoint1/
                  ...
                subendpoint2/
                  ...
            __init__.py

        Returns:
            A formatted string containing:
            - Directory header showing the path being listed (absolute path)
            - When max_depth=1 (default):
                - Directories (trailing slash) listed before files, alphabetically within each type
                - Numbered list (1. 2. 3. ...)
            - When max_depth>1 or max_depth=None:
                - Indentation (2 spaces per level) represents nesting depth
            - Symlink representation:
                - Symlinks are shown as: name -> target (with trailing slash if target is a directory)
                - Broken or unreadable symlinks are shown as: name -> <unreadable>
            - Special files (sockets, devices, FIFOs) are shown as: <special file>
        """
        return await asyncio.to_thread(
            self._list_directory_impl,
            path,
            max_depth,
            show_all,
        )

    class FindFilesResult(TypedDict):
        search_root: str
        total: int
        matches: list[str]

    def _find_files_impl(self,
                         pattern: str,
                         path: str,
                         limit: int,
                         show_all: bool,
                         ) -> FindFilesResult:
        def scan_collect(directory: StdPath) -> list[str]:
            WC_FLAGS = wc_glob.GLOBSTAR | wc_glob.BRACE
            matches = []
            for entry in scantree_bfs(directory, MAX_SCAN_LIMIT, include_hidden, include_gitignored):
                if len(matches) >= limit: break
                if entry.is_symlink(): continue
                relative = StdPath(entry).relative_to(abs_path).as_posix()
                if wc_glob.globmatch(relative, pattern, flags=WC_FLAGS):
                    matches.append(relative)
            return matches

        MAX_SCAN_LIMIT = 200_000
        abs_path = StdPath(self._resolve_path(path))
        include_hidden = show_all
        include_gitignored = show_all

        if not abs_path.exists():
            raise FileNotFoundError(f"Directory not found at {path}")
        if not abs_path.is_dir():
            raise NotADirectoryError(f"Path {path} is not a directory")

        results = scan_collect(abs_path)
        return {
            "search_root": abs_path.as_posix(),
            "total": len(results),
            "matches": results,
        }

    @builtin_tool(validate=True, defaults=BuiltinToolDefaults(auto_approve=True))
    async def find_files(self,
                         pattern: Annotated[str,
                           """
                           A glob pattern to match against file paths, relative to the search root.
                           Supports `**` for recursive matching and `{a,b}` for brace expansion.
                           Pattern examples:
                           - "*.py"          → .py files in the root of the search directory only
                           - "**/*.py"       → .py files in any subdirectory (recursive)
                           - "src/**/*.ts"   → .ts files anywhere under src/
                           - "**/*.{ts,tsx}" → all TypeScript files recursively
                           """],
                         path: Annotated[str,
                           "The path of the directory to search in (relative to the current working directory)."] = ".",
                         limit: Annotated[int,
                           "The maximum number of matching file paths to return."] = 60,
                         show_all: Annotated[bool,
                           "Whether to include hidden files and files ignored by .gitignore. "
                           "Use this if you can't find a specific file you're looking for."] = False
                         ) -> FindFilesResult:
        """
        Search for files whose **NAMES or PATHS** match the glob pattern within a directory.
        Use this when you need to locate files by name, extension, or path structure.

        Note:
            This tool matches file paths/names ONLY. It does NOT read or search inside file contents.

        Examples:
            # Find all Python files in the src directory
            >>> find_files(pattern="**/*.py", path="src")
            {
                "search_root": "/workspace/src",
                "total": 3,
                "matches": [
                    "main.py",
                    "utils/helper.py",
                    "tests/test_main.py"
                ]
            }

            # Find all Markdown and MDX files in the docs directory
            >>> find_files(pattern="docs/*.{md,mdx}")
            {
                "search_root": "/workspace",
                "total": 2,
                "matches": [
                    "docs/README.md",
                    "docs/overview.mdx"
                ]
            }

        Returns:
            A JSON object containing the search results.
            The object has the following properties:
            - search_root: The root directory for the search.
            - total: The total number of matching files found.
            - matches: A list of relative file paths that match the pattern.
        """
        return await asyncio.to_thread(
            self._find_files_impl,
            pattern,
            path,
            limit,
            show_all,
        )

    @builtin_tool(validate=True, defaults=BuiltinToolDefaults(auto_approve=True))
    async def search_text(self,
                          regex: Annotated[str,
                            "The search pattern. Supports regular expressions by default. Use simple strings for literal searches."],
                          path: Annotated[str,
                            "The path to search in (relative to the current working directory), accepts either a directory path or a specific file path."] = ".",
                          file_pattern: Annotated[str | None,
                            """
                            A glob pattern to filter which files are searched, matched against file names and paths.
                            Pattern examples:
                            - "*.py"       → files whose name ends with .py
                            - "*.{ts,tsx}" → TypeScript files
                            - "docs/*.md"  → .md files inside the "docs/" directory
                            - "!*.test.py" → exclude test files
                            Leave unset to search all non-binary files.
                            """
                            ] = None
                         ) -> str:
        """
        Perform a regex search across files in a specified directory, providing context-rich results.
        This tool searches for patterns or specific content across multiple files, displaying each match with encapsulating context.

        Craft your regex patterns carefully to balance specificity and flexibility.
        Use this tool to find any text-based information across the project.
        The results include surrounding context, so analyze the surrounding code to better understand the matches.

        Example:
            >>> search_text(regex="def main\\(", path="src", file_pattern="*.py")
            src/main.py
            52-    return result
            53-
            54:def main():
            55-    parser = argparse.ArgumentParser()
            56-    parser.add_argument("--args")
        """
        resolved_path = str(self._resolve_path(path))
        args = [
            "-n", # show line numbers
            "-C", "2", # 2 lines of context
            "-m", "30", # max 30 matches
            "--path-separator", "/",
            "--smart-case",

            # prevents super long lines (base64, minified js, etc.) and leaves preview and truncated message
            "--max-columns", "300", "--max-columns-preview",

            regex,
            resolved_path,
        ]
        if file_pattern:
            args += ["--glob", file_pattern]
        proc = await asyncio.create_subprocess_exec(
            RIPGREP_PATH,
            *args,
            cwd=self._ctx.cwd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()
        if proc.returncode not in (0, 1):
            raise Exception(stderr.decode())
        return stdout.decode() or "[System] No matches found."
