from __future__ import annotations

from typing import Annotated, Literal

from langchain.tools import InjectedToolCallId, ToolRuntime, tool
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from langgraph.typing import ContextT

from deerflow.agents.thread_state import CanvasArtifact, CanvasArtifactCode, CanvasArtifactText, ThreadState


@tool("update_artifact", parse_docstring=True)
def update_artifact_tool(
    runtime: ToolRuntime[ContextT, ThreadState],
    tool_call_id: Annotated[str, InjectedToolCallId],
    type: Literal["text", "code"],
    title: str,
    content: str,
    language: str = "other",
) -> Command:
    """Create or update the canvas artifact. Always call this to persist your work.

    Args:
        type: Artifact type — "text" for markdown documents, "code" for source code.
        title: Human-readable title of the artifact.
        content: Full artifact content: markdown text or complete source code.
        language: Programming language (for code type only, e.g. "python", "typescript"). Ignored for text.
    """
    current: CanvasArtifact | None = (runtime.state or {}).get("artifact") if runtime.state else None

    if current is not None:
        existing_contents = list(current.get("contents", []))
        new_index = current.get("currentIndex", len(existing_contents)) + 1
    else:
        existing_contents = []
        new_index = 1

    if type == "text":
        new_entry: CanvasArtifactText | CanvasArtifactCode = CanvasArtifactText(
            index=new_index,
            type="text",
            title=title,
            fullMarkdown=content,
        )
    else:
        new_entry = CanvasArtifactCode(
            index=new_index,
            type="code",
            title=title,
            language=language,
            code=content,
        )

    updated: CanvasArtifact = {
        "currentIndex": new_index,
        "contents": existing_contents + [new_entry],
    }

    return Command(
        update={
            "artifact": updated,
            "messages": [ToolMessage("Artifact updated.", tool_call_id=tool_call_id)],
        }
    )