from __future__ import annotations

import asyncio
import logging
import shlex
import subprocess
from abc import ABC, abstractmethod
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any, ClassVar, Self, final, override

import aiofiles
from openai import pydantic_function_tool
from openai.types.chat import ChatCompletionToolParam
from pydantic import BaseModel, Field, ValidationError

from relace_toolbox.async_tasks import AsyncTaskManager, TaskStatus
from relace_toolbox.errors import (
    ExitAgentLoop,
    InvalidToolInput,
    InvalidToolName,
    ToolExecutionError,
    ToolExecutionTimeout,
)
from relace_toolbox.types import ToolResult

logger = logging.getLogger(__name__)

DEFAULT_SUSPEND_TIMEOUT = 2  # seconds


class ToolSchema(BaseModel):
    """Base class for tool input parameters."""

    tool_name: ClassVar[str]
    tool_description: ClassVar[str]

    _tool_lookup: ClassVar[dict[str, type[ToolSchema]]] = {}

    def __init_subclass__(
        cls,
        name: str,
        description: str | None = None,
        register: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init_subclass__(**kwargs)
        description = description or cls.__doc__
        if not description:
            raise ValueError(f"{cls.__name__} is missing a docstring or description")
        cls.tool_name = name
        cls.tool_description = description.strip()
        if register:
            cls._tool_lookup[name] = cls

    @classmethod
    def from_name(cls, tool_name: str) -> type[ToolSchema]:
        try:
            return cls._tool_lookup[tool_name]
        except KeyError as e:
            raise InvalidToolName(tool_name) from e

    @classmethod
    def openai(cls) -> ChatCompletionToolParam:
        schema = pydantic_function_tool(
            model=cls,
            name=cls.tool_name,
            description=cls.tool_description,
        )

        # Remove unsupported 'format' from schema for Path fields, recursively
        def remove_path_format(obj: Any) -> None:
            if isinstance(obj, dict):
                if obj.get("format") == "path":
                    obj.pop("format")
                for val in obj.values():
                    remove_path_format(val)
            elif isinstance(obj, list):
                for item in obj:
                    remove_path_format(item)

        schema_params = schema["function"]["parameters"]
        remove_path_format(schema_params)
        return schema

    @classmethod
    def from_json_str(cls, json_str: str) -> Self:
        try:
            return cls.model_validate_json(json_str)
        except ValidationError as e:
            errors = [f"Error at {err['loc']}: {err['msg']}" for err in e.errors()]
            raise InvalidToolInput("\n".join(("Invalid input:", *errors))) from e

    @classmethod
    def from_json_dict(cls, json_dict: dict[str, Any]) -> Self:
        try:
            return cls.model_validate(json_dict)
        except ValidationError as e:
            errors = [f"Error at {err['loc']}: {err['msg']}" for err in e.errors()]
            raise InvalidToolInput("\n".join(("Invalid input:", *errors))) from e

    def event_data(self) -> dict[str, Any]:
        """Optional event data to include in event stream when calling this tool."""
        return {}


@dataclass
class ToolContext:
    workdir: Path
    chat_history: str | None
    task_manager: AsyncTaskManager


class Tool[S: ToolSchema](ABC):
    """Base class for tools that an agent can use."""

    schema: ClassVar[type[S]]
    requires_chat_history: ClassVar[bool] = False

    _tools: ClassVar[dict[str, type[Tool[ToolSchema]]]] = {}

    def __init_subclass__(
        cls,
        schema: type[S] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init_subclass__(**kwargs)
        # schema is None for abstract base classes like AsyncTool
        if schema is not None:
            cls.schema = schema
            cls._tools[schema.tool_name] = cls  # type: ignore

    @abstractmethod
    async def execute(
        self,
        tool_input: S,
        context: ToolContext,
    ) -> ToolResult:
        """Execute the tool with the given parameters.

        Args:
            tool_input: Structured, validated input arguments for the tool.
        """
        ...

    @classmethod
    def from_name(cls, name: str) -> Tool[ToolSchema]:
        try:
            return cls._tools[name]()
        except KeyError as e:
            raise InvalidToolName(name) from e

    @classmethod
    def all_tools(cls) -> list[type[Tool[ToolSchema]]]:
        return list(cls._tools.values())


class AsyncToolSchema(ToolSchema, name="", register=False):
    """Base class for tool schemas that support async execution with suspension.

    Subclasses must provide their own name and description.
    """

    wait_before_suspend: float = Field(
        DEFAULT_SUSPEND_TIMEOUT,
        description="Maximum seconds to wait before suspending. If the operation "
        "takes longer, returns a task ID to check later.",
    )


class AsyncTool[S: AsyncToolSchema](Tool[S], ABC):
    """Base class for tools that support async execution with suspension.

    Subclasses implement `execute_async()` with the actual tool logic. This base
    class wraps it with suspend/resume logic when `wait_before_suspend` is set.
    """

    cancel_after: ClassVar[float] = 5 * 60  # 5 minutes

    @abstractmethod
    async def execute_async(self, tool_input: S, context: ToolContext) -> ToolResult:
        """Implement the actual tool logic here."""
        ...

    @override
    async def execute(
        self,
        tool_input: S,
        context: ToolContext,
    ) -> ToolResult:
        task_id = await context.task_manager.submit_task(
            coro=self.execute_async(tool_input, context),
            cancel_after=self.cancel_after,
        )
        task_coro = await context.task_manager.wait_for_task(
            task_id=task_id,
            wait_before_suspend=tool_input.wait_before_suspend,
        )

        if task_coro.status == TaskStatus.RUNNING:
            return ToolResult(
                message=(
                    f"Task suspended after {task_coro.elapsed_seconds:.2f}s. "
                    f"Task ID: {task_id}\n"
                    f"Use {CheckTaskSchema.tool_name} to check completion."
                )
            )

        if task_coro.status == TaskStatus.TIMEOUT:
            raise ToolExecutionTimeout(
                f"Operation timed out after {self.cancel_after} seconds."
            )

        if task_coro.status == TaskStatus.FAILED:
            raise ToolExecutionError(f"Operation failed: {task_coro.error}")

        # Completed successfully
        if task_coro.result is None:
            raise ToolExecutionError("Task completed but returned no result")
        return task_coro.result


class BashToolSchema(AsyncToolSchema, name="bash"):
    """Tool for executing bash commands.

    * Avoid dangerous/destructive commands
    * Prefer using other more specialized tools where possible
    * Use wait_before_suspend to limit how long to wait before suspending
    """

    command: str = Field(..., description="Bash command to execute")

    @override
    def event_data(self) -> dict[str, str]:
        return {"command": self.command}


@final
class BashTool(AsyncTool[BashToolSchema], schema=BashToolSchema):
    text_limit: ClassVar[int] = 5_000

    @override
    async def execute_async(
        self,
        tool_input: BashToolSchema,
        context: ToolContext,
    ) -> ToolResult:
        def truncate(text: str) -> str:
            if len(text) > self.text_limit:
                return (
                    f"{text[: self.text_limit]}\n"
                    f"...truncated to {self.text_limit} characters ..."
                )
            return text

        proc = await asyncio.create_subprocess_shell(
            # NOTE: We must run in a clean login shell to prevent the toolbox
            # venv from superceding system python and PYTHONPATH
            f"env -i bash -lc {shlex.quote(tool_input.command)}",
            cwd=context.workdir,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()
        returncode = proc.returncode or 0

        if returncode != 0:
            raise ToolExecutionError(
                f"Command failed with exit code {returncode}.\n"
                f"STDOUT:\n{truncate(stdout.decode())}\n"
                f"STDERR:\n{truncate(stderr.decode())}"
            )
        return ToolResult(message=truncate(stdout.decode()))


class DirectoryViewToolSchema(AsyncToolSchema, name="view_directory"):
    """Tool for viewing the contents of a directory.

    * Lists contents recursively, relative to the input directory
    * Directories are suffixed with a trailing slash '/'
    * In a git repository, respects .gitignore (even when include_hidden is true)
    * Output is limited to the first 250 items

    Example output:
    file1.txt
    file2.txt
    subdir1/
    subdir1/file3.txt
    """

    path: Path = Field(
        ...,
        description="Relative path to a directory, e.g. `src/module/`.",
    )
    include_hidden: bool = Field(
        False,
        description="If true, include hidden files in the output (false by default).",
    )

    @override
    def event_data(self) -> dict[str, str]:
        return {"path": str(self.path)}


class DirectoryViewTool(
    AsyncTool[DirectoryViewToolSchema], schema=DirectoryViewToolSchema
):
    limit: ClassVar[int] = 250

    @override
    async def execute_async(
        self,
        tool_input: DirectoryViewToolSchema,
        context: ToolContext,
    ) -> ToolResult:
        path = resolve_path(
            tool_input.path,
            workdir=context.workdir,
            must_exist=True,
            file_ok=False,
        )

        # Try to get git-tracked files if we're in a git repo
        git_files = await self._get_git_files(path)
        if git_files is not None:
            items = self._iter_git_files(git_files)
        else:
            items = path.rglob("*")

        if not tool_input.include_hidden:
            items = (p for p in items if not p.name.startswith("."))

        contents: list[str] = []
        for i, item in enumerate(items):
            if i >= self.limit:
                contents.append("... remaining items omitted ...")
                break
            relative_path = item.relative_to(path)
            contents.append(
                f"{relative_path}/" if item.is_dir() else str(relative_path)
            )

        return ToolResult(
            message="\n".join(contents) if contents else "Directory is empty."
        )

    async def _get_git_files(self, path: Path) -> set[Path] | None:
        """Get all git-relevant files (tracked + untracked but not ignored).

        Returns None if not in a git repository.
        """
        # Check if we're in a git repo
        proc = await asyncio.create_subprocess_exec(
            "git",
            "rev-parse",
            "--show-toplevel",
            cwd=path,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, _ = await proc.communicate()
        if proc.returncode != 0:
            return None

        # Get tracked files (committed, staged, or modified)
        # and untracked files that are not ignored
        proc = await asyncio.create_subprocess_exec(
            "git",
            "ls-files",
            "--cached",  # tracked files
            "--others",  # untracked files
            "--exclude-standard",  # respect .gitignore
            cwd=path,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, _ = await proc.communicate()
        if proc.returncode != 0:
            return None

        git_files: set[Path] = set()
        for line in stdout.decode().splitlines():
            if line:
                # ls-files outputs paths relative to cwd (which is `path`)
                file_path = path / line
                git_files.add(file_path)
                # Also add all parent directories up to (but not including) path
                for parent in file_path.parents:
                    if parent == path or not parent.is_relative_to(path):
                        break
                    git_files.add(parent)

        return git_files

    def _iter_git_files(self, git_files: set[Path]) -> Iterator[Path]:
        """Iterate existing files from git tracking set, sorted."""
        for item in sorted(git_files):
            if item.exists():
                yield item


class FileViewToolSchema(ToolSchema, name="view_file"):
    """Tool for viewing/exploring the contents of existing files

    Line numbers are included in the output, indexing at 1. If the output does not
    include the end of the file, it will be noted after the final output line.

    Example (viewing the first 2 lines of a file):
    1\tdef my_function():
    2\t    print("Hello, World!")
    ... rest of file truncated ...
    """

    path: Path = Field(
        ...,
        description="Relative path to a file, e.g. `src/file.py`.",
    )
    # TODO: Split this into two integers for more precise validation
    # NOTE: Tuple validation is not supported by the OpenAI API, so we cannot do that
    view_range: list[int] = Field(
        [1, 100],
        description="Range of file lines to view. If not specified, the first 100 lines "
        "of the file are shown. If provided, the file will be shown in the indicated "
        "line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to "
        "start. Setting `[start_line, -1]` shows all lines from `start_line` to the end "
        "of the file.",
    )

    @override
    def event_data(self) -> dict[str, str]:
        return {"path": str(self.path)}


@final
class FileViewTool(Tool[FileViewToolSchema], schema=FileViewToolSchema):
    @override
    async def execute(
        self,
        tool_input: FileViewToolSchema,
        context: ToolContext,
    ) -> ToolResult:
        path = resolve_path(
            tool_input.path,
            workdir=context.workdir,
            must_exist=True,
            dir_ok=False,
        )

        start, end = tool_input.view_range
        view_lines: list[str] = []
        async with aiofiles.open(path) as handle:
            for i, line in enumerate(await handle.readlines(), start=1):
                if end != -1 and i > end:
                    view_lines.append("... rest of file truncated ...")
                    break
                if i >= start:
                    view_lines.append(f"{i}\t{line}")

        return ToolResult(message="".join(view_lines))


class RipgrepToolSchema(AsyncToolSchema, name="grep_search"):
    r"""Fast text-based regex search using ripgrep. Finds exact pattern matches within
    files or directories. Results are capped at 50 matches per file. Best for finding specific
    strings, symbols, or patterns. Preferred over semantic search when you know the
    exact text to find.

    IMPORTANT - Ripgrep uses Rust regex syntax (ERE-style, not BRE):
    - `|`, `+`, `?`, `(`, `)` are special by default. Do NOT escape them like basic grep.
      Use `foo|bar` for alternation, not `foo\|bar`.
    - Lookahead/lookbehind NOT supported
    - Backreferences (`\1`, `\2`) NOT supported

    Use include/exclude patterns to filter by file type or path.
    """

    query: str = Field(..., description="The regex pattern to search for")
    case_sensitive: bool = Field(
        default=True, description="Whether the search should be case sensitive"
    )
    exclude_patterns: list[str] = Field(
        default_factory=list,
        description="Glob patterns for files to exclude",
    )
    include_patterns: list[str] = Field(
        default_factory=list,
        description="Glob patterns for files to include (e.g. '*.ts' for TypeScript files)",
    )

    @override
    def event_data(self) -> dict[str, str]:
        return {"query": self.query}


class RipgrepTool(AsyncTool[RipgrepToolSchema], schema=RipgrepToolSchema):
    @override
    async def execute_async(
        self,
        tool_input: RipgrepToolSchema,
        context: ToolContext,
    ) -> ToolResult:
        args = ["rg", tool_input.query, "./", "--max-count=50", "--color=never"]
        if not tool_input.case_sensitive:
            args.append("--ignore-case")
        for pattern in tool_input.exclude_patterns:
            args.append(f"--glob=!{pattern}")
        for pattern in tool_input.include_patterns:
            args.append(f"--glob={pattern}")

        process = await asyncio.create_subprocess_exec(
            *args,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=context.workdir,
        )
        stdout, stderr = await process.communicate()
        if process.returncode == 0:
            return ToolResult(message=stdout.decode())
        elif process.returncode == 1:
            return ToolResult(message="No matches found")
        else:
            raise ToolExecutionError(f"Error running ripgrep: {stderr.decode()}")


class ReportBackToolSchema(ToolSchema, name="report_back"):
    """This tool is to be used to terminate the run when you have finished exploring
    the codebase and understand the problem.
    """

    explanation: str = Field(
        ...,
        description="Details your reasoning for deeming the files relevant for solving "
        "the issue.",
    )
    files: dict[str, list[tuple[int, int]]] = Field(
        ...,
        description="Key is a relevant file within the repo. The value is the ranges "
        "of relevant lines within the file. Example: {'src/app.py': [(10, 20), (50, 60)]}",
    )

    @override
    def event_data(self) -> dict[str, Any]:
        return {
            "explanation": self.explanation,
            "files": self.files,
        }


class ReportBackTool(Tool[ReportBackToolSchema], schema=ReportBackToolSchema):
    @override
    async def execute(
        self,
        tool_input: ReportBackToolSchema,
        context: ToolContext,
    ) -> ToolResult:
        raise ExitAgentLoop("Report back to user")


class CheckTaskSchema(ToolSchema, name="check_task"):
    """Wait for a previously suspended async task to complete.

    Use this tool to check on or wait for a task that was suspended due to
    wait_before_suspend timeout.
    """

    task_id: str = Field(
        ...,
        description="The task ID returned from a suspended async tool call",
    )
    wait_before_suspend: float = Field(
        DEFAULT_SUSPEND_TIMEOUT,
        description="Maximum seconds to wait before suspending again. "
        "If the task still hasn't completed, returns the same task ID.",
    )


@final
class CheckTaskTool(Tool[CheckTaskSchema], schema=CheckTaskSchema):
    @override
    async def execute(
        self,
        tool_input: CheckTaskSchema,
        context: ToolContext,
    ) -> ToolResult:
        try:
            r = await context.task_manager.wait_for_task(
                task_id=tool_input.task_id,
                wait_before_suspend=tool_input.wait_before_suspend,
            )
        except ValueError as e:
            raise InvalidToolInput(str(e)) from e

        if r.status == TaskStatus.RUNNING:
            return ToolResult(
                message=f"Task still running after {r.elapsed_seconds:.2f}s. "
                f"Call wait_for_task again to continue waiting."
            )

        if r.status == TaskStatus.TIMEOUT:
            return ToolResult(message=f"Task timed out after {r.elapsed_seconds:.2f}s.")

        if r.status == TaskStatus.FAILED:
            return ToolResult(
                message=f"Task failed after {r.elapsed_seconds:.2f}s: {r.error}"
            )

        # Completed successfully
        if r.result is None:
            raise ToolExecutionError("Task completed but returned no result")
        return r.result


def resolve_path(
    path: Path,
    /,
    workdir: Path,
    must_exist: bool = False,
    must_not_exist: bool = False,
    file_ok: bool = True,
    dir_ok: bool = True,
) -> Path:
    if not path.is_absolute():
        path = workdir / path
    elif not path.is_relative_to(workdir):
        raise InvalidToolInput(
            f"path is not contained in the working directory: {path}"
        )
    if must_exist and not path.exists():
        raise InvalidToolInput(f"path does not exist: {path}")
    if must_not_exist and path.exists():
        raise InvalidToolInput(f"path already exists: {path}")
    if not file_ok and path.is_file():
        raise InvalidToolInput(f"path is a file: {path}")
    if not dir_ok and path.is_dir():
        raise InvalidToolInput(f"path is a directory: {path}")
    return path
