"""Runtime-based filesystem for browsing Python module contents."""

from __future__ import annotations

import ast
from dataclasses import dataclass
import inspect
from io import BytesIO
import os
import sys
from types import ModuleType
from typing import Any, Literal, Required, TypedDict, overload

import fsspec

from upathtools.filesystems.base import BaseFileSystem, BaseUPath


NodeType = Literal["function", "class"]


@dataclass
class ModuleMember:
    """A module-level member (function or class)."""

    name: str
    type: NodeType
    doc: str | None = None


class ModuleInfo(TypedDict, total=False):
    """Info dict for module paths and members."""

    name: Required[str]
    type: Required[Literal["module", "function", "class"]]
    size: Required[int]
    doc: Required[str | None]
    mtime: float | None


class ModulePath(BaseUPath[ModuleInfo]):
    """UPath implementation for browsing Python modules."""

    __slots__ = ()

    def iterdir(self):
        if not self.is_dir():
            raise NotADirectoryError(str(self))
        yield from super().iterdir()

    @property
    def path(self) -> str:
        path = super().path
        return "/" if path == "." else path


class ModuleFS(BaseFileSystem[ModulePath, ModuleInfo]):
    """Runtime-based filesystem for browsing a single Python module."""

    protocol = "mod"
    upath_cls = ModulePath

    def __init__(
        self,
        module_path: str = "",
        target_protocol: str | None = None,
        target_options: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)

        # Handle both direct usage and chaining - fo is used by fsspec for chaining
        fo = kwargs.pop("fo", "")
        path = module_path or fo

        if not path:
            msg = "Path to Python file required"
            raise ValueError(msg)

        self.source_path = path if path.endswith(".py") else f"{path}.py"
        self._module: ModuleType | None = None
        self.target_protocol = target_protocol
        self.target_options = target_options or {}

    @staticmethod
    def _get_kwargs_from_urls(path: str) -> dict[str, Any]:
        """Parse mod URL and return constructor kwargs."""
        path = path.removeprefix("mod://")
        return {"module_path": path}

    def _load(self) -> None:
        """Load the module if not already loaded."""
        if self._module is not None:
            return

        # Read and compile the source
        with fsspec.open(
            self.source_path,
            "r",
            protocol=self.target_protocol,
            **self.target_options,
        ) as f:
            source = f.read()  # type: ignore
        self._module = build_module(source, self.source_path)

    @overload
    def ls(self, path: str, detail: Literal[True] = ..., **kwargs: Any) -> list[ModuleInfo]: ...

    @overload
    def ls(self, path: str, detail: Literal[False], **kwargs: Any) -> list[str]: ...

    def ls(self, path: str, detail: bool = True, **kwargs: Any) -> list[ModuleInfo] | list[str]:
        """List module contents (functions and classes)."""
        self._load()
        assert self._module is not None

        members: list[ModuleMember] = []
        for name, obj in vars(self._module).items():
            if name.startswith("_"):
                continue

            if inspect.isfunction(obj):
                member = ModuleMember(name=name, type="function", doc=obj.__doc__)
                members.append(member)
            elif inspect.isclass(obj):
                member = ModuleMember(name=name, type="class", doc=obj.__doc__)
                members.append(member)

        if not detail:
            return [m.name for m in members]

        return [ModuleInfo(name=m.name, type=m.type, doc=m.doc, size=0) for m in members]

    def cat(self, path: str = "") -> bytes:
        """Get source code of whole module or specific member."""
        self._load()
        assert self._module is not None

        path = self._strip_protocol(path).strip("/")  # pyright: ignore[reportAttributeAccessIssue]
        if not path:
            # Return whole module source
            with fsspec.open(
                self.source_path,
                "rb",
                protocol=self.target_protocol,
                **self.target_options,
            ) as f:
                return f.read()  # type: ignore

        # Get specific member
        obj = getattr(self._module, path, None)
        if obj is None:
            msg = f"Member {path} not found"
            raise FileNotFoundError(msg)

        try:
            source = inspect.getsource(obj)
        except OSError:
            # Fallback for Python 3.13+ where inspect.getsource may fail
            source = self._get_source_from_ast(path)
        return source.encode()

    def _get_source_from_ast(self, name: str) -> str:
        """Get source code for a member using AST parsing as fallback."""
        # Read the source file
        with fsspec.open(
            self.source_path,
            "r",
            protocol=self.target_protocol,
            **self.target_options,
        ) as f:
            source_code = f.read()  # pyright: ignore[reportAttributeAccessIssue]

        if (source := get_source_for_node(source_code, name)) is not None:
            return source
        msg = f"Could not find source for {name}"
        raise FileNotFoundError(msg)

    def _open(self, path: str, mode: str = "rb", **kwargs: Any) -> BytesIO:
        """Provide file-like access to source code."""
        if "w" in mode or "a" in mode:
            msg = "Write mode not supported"
            raise NotImplementedError(msg)

        return BytesIO(self.cat(path))

    def info(self, path: str, **kwargs: Any) -> ModuleInfo:
        """Get info about a path."""
        self._load()  # Make sure module is loaded
        assert self._module is not None

        path = self._strip_protocol(path).strip("/")  # pyright: ignore[reportAttributeAccessIssue]

        if not path:
            # Root path - return info about the module itself
            return ModuleInfo(
                name=self._module.__name__,
                type="module",
                size=os.path.getsize(self.source_path),  # noqa: PTH202
                doc=self._module.__doc__,
                mtime=os.path.getmtime(self.source_path)  # noqa: PTH204
                if os.path.exists(self.source_path)  # noqa: PTH110
                else None,
            )

        # Get specific member
        obj = getattr(self._module, path, None)
        if obj is None:
            msg = f"Member {path} not found"
            raise FileNotFoundError(msg)
        return ModuleInfo(
            name=path,
            type="class" if inspect.isclass(obj) else "function",
            size=len(self._get_member_source(obj, path)),
            doc=obj.__doc__,
        )

    def _get_member_source(self, obj: Any, name: str) -> str:
        """Get source code for a member, with fallback for inspect failures."""
        try:
            return inspect.getsource(obj)
        except OSError:
            return self._get_source_from_ast(name)


def build_module(source: str, path: str) -> ModuleType:
    code = compile(source, path, "exec")
    # Create proper module name
    module_name = os.path.splitext(os.path.basename(path))[0]  # noqa: PTH119, PTH122
    # Create module and set up its attributes
    module = ModuleType(module_name)
    module.__file__ = str(path)
    module.__loader__ = None
    module.__package__ = None
    # Register in sys.modules
    sys.modules[module_name] = module
    # Execute in the module's namespace
    exec(code, module.__dict__)
    # Set __module__ for all classes and functions
    for obj in module.__dict__.values():
        if inspect.isclass(obj) or inspect.isfunction(obj):
            obj.__module__ = module_name
    return module


def get_source_for_node(source_code: str, name: str) -> str | None:
    tree = ast.parse(source_code)
    for node in ast.walk(tree):  # Find the node with the given name
        if (
            isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.AsyncFunctionDef))
            and node.name == name
        ):
            # Extract the source lines for this node
            lines = source_code.splitlines()
            start_line = node.lineno - 1
            # Find the end line by looking at indentation
            end_line = len(lines)
            base_indent = len(lines[start_line]) - len(lines[start_line].lstrip())
            for i in range(start_line + 1, len(lines)):
                line = lines[i]
                if line.strip():  # Skip empty lines
                    current_indent = len(line) - len(line.lstrip())
                    if current_indent <= base_indent:
                        end_line = i
                        break

            return "\n".join(lines[start_line:end_line])
    return None


if __name__ == "__main__":
    fs = fsspec.filesystem("mod", module_path="src/upathtools/helpers.py")
    print(fs.info("/"))
    # print(fs.cat("build"))
