from pathlib import Path
from typing import Optional, Callable, Dict
from dataclasses import dataclass
import subprocess
from llm import AgentConfig, JSON


@dataclass
class Action:
    """Class for LLM actions.

    Should only be called on JSON dicts where "action" -> `name`.
    """

    name: str
    description: str
    handler: Callable[[JSON, AgentConfig], JSON]
    formatter: Callable[[JSON], str] = lambda args: (
        f"{args['action']} with "
        + ",".join([f"{arg} as {args[arg]}" for arg in args if arg != "action"])
    )

    def __call__(self, data: JSON, config: AgentConfig) -> JSON:
        return self.handler(data, config)


def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]:
    """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`."""
    path = (path if path.is_absolute() else (cwd / path)).resolve()
    cwd = cwd.resolve()
    return path if path.is_relative_to(cwd) else None


def _path_helper(path_str: str, callback, cwd=Path.cwd()):
    path = _path_checked(Path(path_str), cwd=cwd)
    if path is None:
        return {"status": "error", "reason": f"Permission denied: {path_str}"}
    else:
        try:
            return {"status": "success", **callback(path)}
        except Exception as e:
            return {"status": "error", "reason": f"Failure: {e}"}


def _ls(args: JSON, config: AgentConfig) -> JSON:
    '{"action": "ls", "path": PATH} - list files at PATH'

    def inner(path: Path):
        return {"content": list(map(lambda s: str(s), path.iterdir()))}

    return _path_helper(args["path"], inner, cwd=config.workdir)


def _read(args: JSON, config: AgentConfig) -> JSON:
    '{"action": "read", "path": PATH} - read file at PATH'

    def inner(path: Path):
        contents = path.read_text(encoding="utf-8")
        return {"content": contents}

    return _path_helper(args["path"], inner, cwd=config.workdir)


def _write(args: JSON, config: AgentConfig) -> JSON:
    '{"action": "write", "path": PATH, "content": CONTENT} - write CONTENT into file at PATH'

    def inner(path: Path, text: str):
        n_bytes = path.write_text(text)
        return {"bytes": n_bytes}

    return _path_helper(
        args["path"], lambda path: inner(path, args["content"]), cwd=config.workdir
    )


def _tldr(args: JSON, config: AgentConfig) -> JSON:
    '{"action": "tldr", "program": PROGRAM} - read tldr docs for PROGRAM'
    try:
        result = subprocess.run(
            ["tldr", args["program"]], capture_output=True, text=True, check=True
        )
        return {"status": "success", "content": result.stdout}
    except subprocess.CalledProcessError:
        return {"status": "error", "reason": f"Tldr for '{args['program']}' not found."}
    except Exception as e:
        return {"status": "error", "reason": f"An unexpected error occurred: {e}"}


def _shell(args: JSON, config: AgentConfig) -> JSON:
    '{"action": "shell", "command": CMD} - Propose executing a shell command, but interrupts to ask for permission from the user.'
    res = subprocess.run(
        args["command"],
        shell=True,
        cwd=config.workdir,
        capture_output=True,
        text=True,
    )
    if res.returncode == 0:
        return {"status": "success", "stdout": res.stdout}
    else:
        return {
            "status": "error",
            "reason": f"Shell command exited with code {res.returncode}",
            "stdout": res.stdout,
            "stderr": res.stderr,
        }


ls = Action("ls", _ls.__doc__ or "", _ls)
read = Action("read", _read.__doc__ or "", _read)
write = Action(
    "write",
    _write.__doc__ or "",
    _write,
    formatter=lambda args: f"write {args['path']}",
)
tldr = Action("tldr", _tldr.__doc__ or "", _tldr)
shell = Action(
    "shell",
    _shell.__doc__ or "",
    _shell,
    formatter=lambda args: f"shell command `{args['command']}`",
)