"""Session-scoped scratchpad for capturing working notes during task execution — backed by PostgreSQL KV store."""
from __future__ import annotations

from navi.tools._internal.base import Tool, ToolResult, current_session_id, current_user_id

# Global KV store reference — injected at startup by registry.py
_kv_store = None


def set_kv_store(kv) -> None:
    """Inject the shared KvStore instance (called once at startup)."""
    global _kv_store
    _kv_store = kv


def _sid() -> str:
    return current_session_id.get() or "__default__"


def _uid() -> str | None:
    return current_user_id.get(None)


async def get_section(session_id: str, section: str) -> str:
    """Read one scratchpad section for the given session. Returns '' if absent."""
    if _kv_store is None:
        return ""
    try:
        val = await _kv_store.get(None, session_id, "scratchpad", section)
        return val or ""
    except Exception:
        return ""


class ScratchpadTool(Tool):
    name = "scratchpad"
    description = (
        "Working memory for the current session — for facts discovered mid-task, not for progress tracking. "
        "Sections: goal (one-sentence objective), findings (key facts from tool results), "
        "artifacts (file paths, URLs), errors (failures and what was tried). "
        "Write 'goal' at the start of any multi-step task. "
        "Read before composing a final answer — findings may contain facts needed for the response."
    )
    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. Standard names: "
                    "goal, findings, artifacts, errors, main. "
                    "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"],
    }

    def __init__(self, kv_store=None) -> None:
        if kv_store is not None:
            set_kv_store(kv_store)

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

        if op == "write":
            if not content:
                return ToolResult(success=False, output="", error="'content' is required for 'write'")
            key = section or "main"
            if _kv_store is not None:
                await _kv_store.set(_uid(), sid, "scratchpad", 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"
            if _kv_store is not None:
                existing = await _kv_store.get(_uid(), sid, "scratchpad", key) or ""
                new = (existing + "\n" + content).lstrip("\n") if existing else content
                await _kv_store.set(_uid(), sid, "scratchpad", key, new)
                return ToolResult(success=True, output=f"[{key}] updated ({len(new)} chars total).")
            return ToolResult(success=True, output=f"[{key}] updated.")

        if op == "read":
            if _kv_store is None:
                return ToolResult(success=True, output="Scratchpad is empty.")
            if section is not None:
                text = await _kv_store.get(_uid(), sid, "scratchpad", 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
            all_data = await _kv_store.get_all(_uid(), sid, "scratchpad")
            if not all_data:
                return ToolResult(success=True, output="Scratchpad is empty.")
            parts = [f"[{k}]:\n{v}" for k, v in all_data.items()]
            return ToolResult(success=True, output="\n\n".join(parts))

        if op == "clear":
            if _kv_store is None:
                return ToolResult(success=True, output="Scratchpad cleared.")
            if section is not None:
                existing = await _kv_store.get(_uid(), sid, "scratchpad", section)
                await _kv_store.delete(_uid(), sid, "scratchpad", section)
                return ToolResult(
                    success=True,
                    output=f"[{section}] cleared." if existing else f"[{section}] was already empty.",
                )
            await _kv_store.clear_scope(_uid(), sid, "scratchpad")
            return ToolResult(success=True, output="Scratchpad cleared.")

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