"""Session-scoped scratchpad for capturing working notes during task execution."""
from __future__ import annotations

from .base import Tool, ToolResult, current_session_id

# session_id → {section_name: content}
_pads: dict[str, dict[str, str]] = {}


def get_section(session_id: str, section: str) -> str:
    """Read one scratchpad section for the given session. Returns '' if absent."""
    return _pads.get(session_id, {}).get(section, "")


class ScratchpadTool(Tool):
    name = "scratchpad"
    description = (
        "Persistent working memory for the current session. Use it to record findings, "
        "track artifacts, and pass context to sub-agents. Organised into named sections. "
        "MANDATORY: at the start of any multi-step task, write a 'goal' section stating "
        "in one sentence what you are trying to achieve. Update it if the goal changes. "
        "Always read the relevant section before composing a final answer — "
        "it may contain findings from earlier tool calls."
    )
    parameters = {
        "type": "object",
        "properties": {
            "op": {
                "type": "string",
                "enum": ["write", "append", "read", "clear"],
                "description": (
                    "write  — overwrite a section with new content (requires 'content'). "
                    "append — add text to the end of an existing section (requires 'content'). "
                    "read   — return one section (requires 'section') or all sections (omit 'section'). "
                    "clear  — erase one section (requires 'section') or the entire pad (omit 'section')."
                ),
            },
            "section": {
                "type": "string",
                "description": (
                    "Which section to target. Conventional names: "
                    "goal (current objective — always write this first), "
                    "findings (discovered facts), artifacts (file paths, outputs), "
                    "errors (tracebacks, failures), verification (DoD checklist), "
                    "context_transfer (briefings for sub-agents), main (general). "
                    "Defaults to 'main' when omitted on write/append. "
                    "Omit entirely on read/clear to target all sections."
                ),
            },
            "content": {
                "type": "string",
                "description": (
                    "Text to write or append. REQUIRED for 'write' and 'append' — "
                    "the call will fail without it."
                ),
            },
        },
        "required": ["op"],
    }

    async def execute(self, params: dict) -> ToolResult:
        sid = current_session_id.get() or "__default__"
        op = params.get("op")
        section: str | None = params.get("section") or None
        content: str = params.get("content", "")

        pad = _pads.setdefault(sid, {})

        if op == "write":
            if not content:
                return ToolResult(success=False, output="", error="'content' is required for 'write'")
            key = section or "main"
            pad[key] = content
            return ToolResult(success=True, output=f"[{key}] written ({len(content)} chars).")

        if op == "append":
            if not content:
                return ToolResult(success=False, output="", error="'content' is required for 'append'")
            key = section or "main"
            existing = pad.get(key, "")
            pad[key] = (existing + "\n" + content).lstrip("\n") if existing else content
            return ToolResult(success=True, output=f"[{key}] updated ({len(pad[key])} chars total).")

        if op == "read":
            if section is not None:
                text = pad.get(section)
                if not text:
                    return ToolResult(success=True, output=f"[{section}] is empty.")
                return ToolResult(success=True, output=f"[{section}]:\n{text}")
            # No section → read all
            if not pad:
                return ToolResult(success=True, output="Scratchpad is empty.")
            parts = [f"[{k}]:\n{v}" for k, v in pad.items()]
            return ToolResult(success=True, output="\n\n".join(parts))

        if op == "clear":
            if section is not None:
                removed = pad.pop(section, None)
                return ToolResult(
                    success=True,
                    output=f"[{section}] cleared." if removed else f"[{section}] was already empty.",
                )
            pad.clear()
            return ToolResult(success=True, output="Scratchpad cleared.")

        return ToolResult(success=False, output="", error=f"Unknown op: {op!r}")