from typing import Annotated, Literal, NotRequired, Optional, TypedDict

from langchain.agents import AgentState


class SandboxState(TypedDict):
    sandbox_id: NotRequired[str | None]


class ThreadDataState(TypedDict):
    workspace_path: NotRequired[str | None]
    uploads_path: NotRequired[str | None]
    outputs_path: NotRequired[str | None]


class ViewedImageData(TypedDict):
    base64: str
    mime_type: str


class CanvasArtifactText(TypedDict):
    index: int
    type: Literal["text"]
    title: str
    fullMarkdown: str


class CanvasArtifactCode(TypedDict):
    index: int
    type: Literal["code"]
    title: str
    language: str
    code: str


class CanvasArtifact(TypedDict):
    currentIndex: int
    contents: list[CanvasArtifactText | CanvasArtifactCode]


def merge_artifacts(existing: list[str] | None, new: list[str] | None) -> list[str]:
    """Reducer for artifacts list - merges and deduplicates artifacts."""
    if existing is None:
        return new or []
    if new is None:
        return existing
    # Use dict.fromkeys to deduplicate while preserving order
    return list(dict.fromkeys(existing + new))


def merge_viewed_images(existing: dict[str, ViewedImageData] | None, new: dict[str, ViewedImageData] | None) -> dict[str, ViewedImageData]:
    """Reducer for viewed_images dict - merges image dictionaries.

    Special case: If new is an empty dict {}, it clears the existing images.
    This allows middlewares to clear the viewed_images state after processing.
    """
    if existing is None:
        return new or {}
    if new is None:
        return existing
    # Special case: empty dict means clear all viewed images
    if len(new) == 0:
        return {}
    # Merge dictionaries, new values override existing ones for same keys
    return {**existing, **new}


class ThreadState(AgentState):
    sandbox: NotRequired[SandboxState | None]
    thread_data: NotRequired[ThreadDataState | None]
    title: NotRequired[str | None]
    # True while `title` is a provisional, question-derived placeholder that the
    # async LLM pass has not yet finalized. Lets TitleMiddleware show a title
    # instantly (no "untitled" window) and upgrade it once the answer is ready.
    title_provisional: NotRequired[bool]
    artifacts: Annotated[list[str], merge_artifacts]
    todos: NotRequired[list | None]
    uploaded_files: NotRequired[list[dict] | None]
    viewed_images: Annotated[dict[str, ViewedImageData], merge_viewed_images]  # image_path -> {base64, mime_type}
    # Canvas fields (coexist with artifacts: list[str] which tracks file outputs)
    artifact: NotRequired[Optional[CanvasArtifact]]
    highlighted_text: NotRequired[Optional[str]]
    highlighted_code: NotRequired[Optional[str]]
    canvas_mode: NotRequired[Optional[Literal["text", "code"]]]
    quick_action_id: NotRequired[Optional[str]]
    web_search_enabled: NotRequired[bool]
    web_search_results: NotRequired[Optional[list[dict]]]
    context_summary_cache: NotRequired[dict | None]