# Copyright (c) 2023-2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import asyncio
import codecs
import contextlib
import inspect
import json
import logging
import shutil
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass
from dataclasses import field as datafield
from datetime import datetime, timezone
from functools import partial
from pathlib import Path
from textwrap import shorten
from typing import Any, Literal

from pydantic import BaseModel, Field, field_validator, model_validator

from lionagi import ln
from lionagi.libs.schema.as_readable import as_readable

HAS_CLAUDE_CODE_CLI = False
CLAUDE_CLI = None

if (c := (shutil.which("claude") or "claude")) and shutil.which(c):
    HAS_CLAUDE_CODE_CLI = True
    CLAUDE_CLI = c

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("claude-cli")


# --------------------------------------------------------------------------- types
ClaudePermissionMode = Literal[
    "default",
    "acceptEdits",
    "plan",
    "dontAsk",
    "bypassPermissions",
]
# Backward-compat alias
ClaudePermission = ClaudePermissionMode

ClaudeEffort = Literal["low", "medium", "high", "max"]
ClaudeOutputFormat = Literal["text", "json", "stream-json"]
ClaudeInputFormat = Literal["text", "stream-json"]


__all__ = (
    "ClaudeCodeRequest",
    "ClaudeChunk",
    "ClaudeSession",
    "stream_claude_code_cli",
)


# --------------------------------------------------------------------------- flag metadata
#
# Each CLI-mappable field carries a ``json_schema_extra`` dict produced by
# ``_cli()``.  The generic ``_build_declarative_args()`` loop reads these
# dicts (sorted by *order*) and emits the correct flag sequence.
#
# kind semantics:
#   value      – ``--flag <str(val)>``
#   bool       – ``--flag`` when truthy, omit otherwise
#   bool_pair  – ``--flag`` when True, ``--neg-flag`` when False, omit when None
#   list_args  – ``--flag arg1 arg2 …`` (one flag, many positional args)
#   json_value – ``--flag '<json>'``   (dict/list serialised to JSON string)
#   repeat     – ``--flag a --flag b`` (flag repeated per item)


def _cli(
    flag: str,
    order: int,
    kind: str = "value",
    *,
    neg_flag: str | None = None,
) -> dict[str, Any]:
    d: dict[str, Any] = {
        "cli_flag": flag,
        "cli_order": order,
        "cli_kind": kind,
    }
    if neg_flag:
        d["cli_neg_flag"] = neg_flag
    return d


# --------------------------------------------------------------------------- request model
class ClaudeCodeRequest(BaseModel):
    """Configuration + prompt for a Claude Code CLI invocation.

    Every field annotated with ``_cli(...)`` metadata is automatically
    assembled into CLI arguments by :meth:`as_cmd_args`, sorted by its
    ``order`` value.  A handful of special cases (``permission_mode``,
    ``max_turns``, ``worktree``, ``debug``, legacy ``mcp_servers``) are
    handled with explicit logic after the declarative pass.

    Adding a new CLI flag is one line: declare the field with ``_cli()``
    metadata and the builder picks it up automatically.
    """

    # ── prompt (always required) ──────────────────────────────────
    prompt: str = Field(description="The prompt for Claude Code")

    # ── session management (order 10–19) ──────────────────────────
    continue_conversation: bool = Field(
        default=False,
        json_schema_extra=_cli("--continue", 10, "bool"),
    )
    resume: str | None = Field(
        default=None,
        json_schema_extra=_cli("--resume", 11),
    )
    session_id: str | None = Field(
        default=None,
        json_schema_extra=_cli("--session-id", 12),
    )
    name: str | None = Field(
        default=None,
        json_schema_extra=_cli("--name", 13),
    )
    fork_session: bool = Field(
        default=False,
        json_schema_extra=_cli("--fork-session", 14, "bool"),
    )

    # ── model & runtime (order 20–29) ─────────────────────────────
    model: Literal["sonnet", "opus", "haiku"] | str | None = Field(
        default="sonnet",
        json_schema_extra=_cli("--model", 20),
    )
    effort: ClaudeEffort | None = Field(
        default=None,
        json_schema_extra=_cli("--effort", 21),
    )
    fallback_model: str | None = Field(
        default=None,
        json_schema_extra=_cli("--fallback-model", 22),
    )
    # max_turns is special-cased (+1 offset) — no _cli metadata
    max_turns: int | None = None
    max_budget_usd: float | None = Field(
        default=None,
        json_schema_extra=_cli("--max-budget-usd", 24),
    )

    # ── system prompt (order 30–39) ───────────────────────────────
    system_prompt: str | None = Field(
        default=None,
        json_schema_extra=_cli("--system-prompt", 30),
    )
    system_prompt_file: str | Path | None = Field(
        default=None,
        json_schema_extra=_cli("--system-prompt-file", 31),
    )
    append_system_prompt: str | None = Field(
        default=None,
        json_schema_extra=_cli("--append-system-prompt", 32),
    )
    append_system_prompt_file: str | Path | None = Field(
        default=None,
        json_schema_extra=_cli("--append-system-prompt-file", 33),
    )

    # ── permissions (order 40–49) ─────────────────────────────────
    # permission_mode is special-cased — no _cli metadata
    permission_mode: ClaudePermissionMode | None = None
    allow_dangerously_skip_permissions: bool = Field(
        default=False,
        json_schema_extra=_cli("--allow-dangerously-skip-permissions", 40, "bool"),
    )
    allowed_tools: list[str] | None = Field(
        default=None,
        json_schema_extra=_cli("--allowedTools", 42, "list_args"),
    )
    disallowed_tools: list[str] | None = Field(
        default=None,
        json_schema_extra=_cli("--disallowedTools", 43, "list_args"),
    )
    tools: str | None = Field(
        default=None,
        description="Restrict available tools (comma-separated or 'default')",
        json_schema_extra=_cli("--tools", 44),
    )
    permission_prompt_tool_name: str | None = Field(
        default=None,
        json_schema_extra=_cli("--permission-prompt-tool", 45),
    )

    # ── MCP (order 50–59) ─────────────────────────────────────────
    mcp_config: str | Path | None = Field(
        default=None,
        json_schema_extra=_cli("--mcp-config", 50),
    )
    strict_mcp_config: bool = Field(
        default=False,
        json_schema_extra=_cli("--strict-mcp-config", 51, "bool"),
    )
    # Legacy: if set and mcp_config is absent, serialised to --mcp-config JSON
    mcp_servers: dict[str, Any] = Field(default_factory=dict, exclude=True)

    # ── agents (order 60–69) ──────────────────────────────────────
    agent: str | None = Field(
        default=None,
        json_schema_extra=_cli("--agent", 60),
    )
    agents: dict[str, Any] | None = Field(
        default=None,
        json_schema_extra=_cli("--agents", 61, "json_value"),
    )

    # ── workspace (order 70–79) ───────────────────────────────────
    repo: Path = Field(default_factory=Path.cwd, exclude=True)
    ws: str | None = Field(default=None, exclude=True)
    add_dir: list[str] | None = Field(
        default=None,
        json_schema_extra=_cli("--add-dir", 70, "list_args"),
    )
    # worktree is special-cased (bool vs str) — no _cli metadata
    worktree: str | bool | None = Field(default=None, exclude=True)

    # ── features (order 80–89) ────────────────────────────────────
    chrome: bool | None = Field(
        default=None,
        json_schema_extra=_cli("--chrome", 80, "bool_pair", neg_flag="--no-chrome"),
    )
    disable_slash_commands: bool = Field(
        default=False,
        json_schema_extra=_cli("--disable-slash-commands", 81, "bool"),
    )
    no_session_persistence: bool = Field(
        default=False,
        json_schema_extra=_cli("--no-session-persistence", 82, "bool"),
    )

    # ── output (order 90–99) ──────────────────────────────────────
    output_format: ClaudeOutputFormat = Field(default="stream-json", exclude=True)
    input_format: ClaudeInputFormat | None = Field(
        default=None,
        json_schema_extra=_cli("--input-format", 91),
    )
    # Named json_schema_output to avoid shadowing Pydantic's json_schema
    json_schema_output: dict[str, Any] | str | None = Field(
        default=None,
        json_schema_extra=_cli("--json-schema", 92, "json_value"),
    )
    include_partial_messages: bool = Field(
        default=False,
        json_schema_extra=_cli("--include-partial-messages", 93, "bool"),
    )

    # ── settings & debug (order 100–109) ──────────────────────────
    settings: str | Path | None = Field(
        default=None,
        json_schema_extra=_cli("--settings", 100),
    )
    setting_sources: str | None = Field(
        default=None,
        json_schema_extra=_cli("--setting-sources", 101),
    )
    betas: list[str] | None = Field(
        default=None,
        json_schema_extra=_cli("--betas", 102, "list_args"),
    )
    # debug is special-cased (bool vs str) — no _cli metadata
    debug: str | bool | None = Field(default=None, exclude=True)

    # ── lionagi internal (never become CLI flags) ─────────────────
    auto_finish: bool = Field(
        default=False,
        exclude=True,
        description="Automatically finish the conversation after the first response",
    )
    verbose_output: bool = Field(default=False, exclude=True)
    cli_display_theme: Literal["light", "dark"] = Field(default="dark", exclude=True)
    cli_include_summary: bool = Field(default=False, exclude=True)

    # Legacy fields (kept for backward compat, unused in CLI args)
    mcp_tools: list[str] = Field(default_factory=list, exclude=True)
    max_thinking_tokens: int | None = Field(default=None, exclude=True)

    # ── validators ────────────────────────────────────────────────

    @field_validator("permission_mode", mode="before")
    def _norm_perm(cls, v):
        if v in {
            "dangerously-skip-permissions",
            "--dangerously-skip-permissions",
        }:
            return "bypassPermissions"
        return v

    @field_validator("add_dir", mode="before")
    def _norm_add_dir(cls, v):
        if isinstance(v, str):
            return [v]
        return v

    @model_validator(mode="before")
    def _validate_message_prompt(cls, data):
        if "prompt" in data and data["prompt"]:
            return data

        if not (msg := data.get("messages")):
            raise ValueError("messages may not be empty")
        resume = data.get("resume")
        continue_conversation = data.get("continue_conversation")

        prompt = ""

        # 1. if resume or continue_conversation, use the last message
        if resume or continue_conversation:
            continue_conversation = True
            prompt = msg[-1]["content"]
            if isinstance(prompt, (dict, list)):
                prompt = ln.json_dumps(prompt)

        # 2. else, use entire messages except system message
        else:
            prompts = []
            continue_conversation = False
            for message in msg:
                if message["role"] != "system":
                    content = message["content"]
                    prompts.append(
                        ln.json_dumps(content)
                        if isinstance(content, (dict, list))
                        else content
                    )

            prompt = "\n".join(prompts)

        # 3. assemble the request data
        data_: dict[str, Any] = dict(
            prompt=prompt,
            resume=resume,
            continue_conversation=bool(continue_conversation),
        )

        # 4. extract system prompt if available
        if msg[0]["role"] == "system":
            if resume or continue_conversation:
                # Continued session: pass as system_prompt (--system-prompt)
                data_["system_prompt"] = msg[0]["content"]
            else:
                # Fresh session: append to Claude Code's built-in prompt
                data_.setdefault("append_system_prompt", msg[0]["content"])

        if "append_system_prompt" in data and data["append_system_prompt"]:
            data_["append_system_prompt"] = str(data.get("append_system_prompt"))

        data_.update(data)
        return data_

    @model_validator(mode="after")
    def _check_constraints(self):
        # Session flag constraints
        if self.resume:
            self.continue_conversation = False
        if self.fork_session and not (self.resume or self.continue_conversation):
            raise ValueError("--fork-session requires --resume or --continue")

        # Permission flag mutual exclusivity
        if (
            self.allow_dangerously_skip_permissions
            and self.permission_mode == "bypassPermissions"
        ):
            raise ValueError(
                "allow_dangerously_skip_permissions and "
                "permission_mode='bypassPermissions' are mutually exclusive"
            )

        # System prompt mutual exclusivity
        if self.system_prompt and self.system_prompt_file:
            raise ValueError(
                "--system-prompt and --system-prompt-file are mutually exclusive"
            )

        # Workspace bounds check for bypassPermissions
        if self.permission_mode == "bypassPermissions":
            repo_resolved = self.repo.resolve()
            cwd_resolved = self.cwd().resolve()
            try:
                cwd_resolved.relative_to(repo_resolved)
            except ValueError:
                raise ValueError(
                    f"With bypassPermissions, workspace must be within "
                    f"repository bounds. Repository: {repo_resolved}, "
                    f"Workspace: {cwd_resolved}"
                ) from None

        return self

    # ── workspace path ────────────────────────────────────────────

    def cwd(self) -> Path:
        if not self.ws:
            return self.repo

        ws_path = Path(self.ws)

        if ws_path.is_absolute():
            raise ValueError(
                f"Workspace path must be relative, got absolute: {self.ws}"
            )

        if ".." in ws_path.parts:
            raise ValueError(
                f"Directory traversal detected in workspace path: {self.ws}"
            )

        repo_resolved = self.repo.resolve()
        result = (self.repo / ws_path).resolve()

        try:
            result.relative_to(repo_resolved)
        except ValueError:
            raise ValueError(
                f"Workspace path escapes repository bounds. "
                f"Repository: {repo_resolved}, Workspace: {result}"
            ) from None

        return result

    # ── CLI command builder ───────────────────────────────────────

    def as_cmd_args(self) -> list[str]:
        """Build argument list for the Claude Code CLI.

        Flags are assembled in two passes:

        1. **Declarative** – fields carrying ``_cli()`` metadata are
           collected, sorted by ``order``, and emitted by ``kind``.
        2. **Special cases** – ``permission_mode``, ``max_turns``,
           ``worktree``, ``debug``, and legacy ``mcp_servers`` need
           non-mechanical logic and are handled explicitly.

        The prompt (``-p``) and ``--output-format`` always come first;
        ``--verbose`` is always appended last.
        """
        args: list[str] = [
            "-p",
            self.prompt,
            "--output-format",
            self.output_format,
        ]

        # ── pass 1: declarative flags ──
        args.extend(self._build_declarative_args())

        # ── pass 2: special cases ──

        # permission_mode → --dangerously-skip-permissions OR --permission-mode
        if self.permission_mode:
            if self.permission_mode == "bypassPermissions":
                args.append("--dangerously-skip-permissions")
            else:
                args.extend(["--permission-mode", self.permission_mode])

        # max_turns → +1 offset (CLI counts agentic turns differently)
        if self.max_turns is not None:
            args.extend(["--max-turns", str(self.max_turns + 1)])

        # worktree → True emits bare flag, str emits flag + name
        if self.worktree is not None and self.worktree is not False:
            if isinstance(self.worktree, str):
                args.extend(["--worktree", self.worktree])
            else:
                args.append("--worktree")

        # debug → True emits bare flag, str emits flag + categories
        if self.debug:
            if isinstance(self.debug, str):
                args.extend(["--debug", self.debug])
            else:
                args.append("--debug")

        # Legacy mcp_servers dict → serialise as --mcp-config JSON inline
        if self.mcp_servers and not self.mcp_config:
            args.extend(
                [
                    "--mcp-config",
                    json.dumps({"mcpServers": self.mcp_servers}),
                ]
            )

        # model default – always emit
        if "--model" not in args:
            args.extend(["--model", self.model or "sonnet"])

        # always verbose for structured stream output
        args.append("--verbose")
        return args

    def _build_declarative_args(self) -> list[str]:
        """Collect fields with ``_cli()`` metadata and emit flags."""
        flagged: list[tuple[int, dict, Any]] = []
        for field_name, field_info in type(self).model_fields.items():
            extra = field_info.json_schema_extra
            if not extra or "cli_flag" not in extra:
                continue
            val = getattr(self, field_name)
            if val is None:
                continue
            if isinstance(val, list) and not val:
                continue
            if val is False and extra.get("cli_kind") != "bool_pair":
                continue
            flagged.append((extra["cli_order"], extra, val))

        flagged.sort(key=lambda x: x[0])

        args: list[str] = []
        for _, extra, val in flagged:
            flag = extra["cli_flag"]
            kind = extra.get("cli_kind", "value")

            if kind == "bool":
                if val:
                    args.append(flag)

            elif kind == "bool_pair":
                if val is True:
                    args.append(flag)
                elif val is False and extra.get("cli_neg_flag"):
                    args.append(extra["cli_neg_flag"])

            elif kind == "list_args":
                args.append(flag)
                args.extend(str(v) for v in val)

            elif kind == "json_value":
                serialized = (
                    json.dumps(val) if isinstance(val, (dict, list)) else str(val)
                )
                args.extend([flag, serialized])

            elif kind == "repeat":
                for v in val:
                    args.extend([flag, str(v)])

            else:  # "value"
                args.extend([flag, str(val)])

        return args


# --------------------------------------------------------------------------- chunks & session


@dataclass
class ClaudeChunk:
    """Low-level wrapper around every NDJSON object coming from the CLI."""

    raw: dict[str, Any]
    type: str
    # convenience views
    thinking: str | None = None
    text: str | None = None
    tool_use: dict[str, Any] | None = None
    tool_result: dict[str, Any] | None = None


@dataclass
class ClaudeSession:
    """Aggregated view of a whole CLI conversation."""

    session_id: str | None = None
    model: str | None = None

    # chronological log
    chunks: list[ClaudeChunk] = datafield(default_factory=list)

    # materialised views
    thinking_log: list[str] = datafield(default_factory=list)
    messages: list[dict[str, Any]] = datafield(default_factory=list)
    tool_uses: list[dict[str, Any]] = datafield(default_factory=list)
    tool_results: list[dict[str, Any]] = datafield(default_factory=list)

    # final summary
    result: str = ""
    usage: dict[str, Any] = datafield(default_factory=dict)
    total_cost_usd: float | None = None
    num_turns: int | None = None
    duration_ms: int | None = None
    duration_api_ms: int | None = None
    is_error: bool = False
    summary: dict | None = None

    def populate_summary(self) -> None:
        self.summary = _extract_summary(self)


def _extract_summary(session: ClaudeSession) -> dict[str, Any]:
    tool_counts = {}
    tool_details = []
    file_operations = {"reads": [], "writes": [], "edits": []}
    key_actions = []

    # Process tool uses from the clean materialized view
    for tool_use in session.tool_uses:
        tool_name = tool_use.get("name", "unknown")
        tool_input = tool_use.get("input", {})
        tool_id = tool_use.get("id", "")

        # Count tool usage
        tool_counts[tool_name] = tool_counts.get(tool_name, 0) + 1

        # Store detailed info
        tool_details.append({"tool": tool_name, "id": tool_id, "input": tool_input})

        # Categorize file operations and actions
        if tool_name in ["Read", "read"]:
            file_path = tool_input.get("file_path", "unknown")
            file_operations["reads"].append(file_path)
            key_actions.append(f"Read {file_path}")

        elif tool_name in ["Write", "write"]:
            file_path = tool_input.get("file_path", "unknown")
            file_operations["writes"].append(file_path)
            key_actions.append(f"Wrote {file_path}")

        elif tool_name in ["Edit", "edit", "MultiEdit"]:
            file_path = tool_input.get("file_path", "unknown")
            file_operations["edits"].append(file_path)
            key_actions.append(f"Edited {file_path}")

        elif tool_name in ["Bash", "bash"]:
            command = tool_input.get("command", "")
            command_summary = command[:50] + "..." if len(command) > 50 else command
            key_actions.append(f"Ran: {command_summary}")

        elif tool_name in ["Glob", "glob"]:
            pattern = tool_input.get("pattern", "")
            key_actions.append(f"Searched files: {pattern}")

        elif tool_name in ["Grep", "grep"]:
            pattern = tool_input.get("pattern", "")
            key_actions.append(f"Searched content: {pattern}")

        elif tool_name in ["Task", "task", "Agent"]:
            description = tool_input.get("description", "")
            key_actions.append(f"Spawned agent: {description}")

        elif tool_name.startswith("mcp__"):
            operation = tool_name.replace("mcp__", "")
            key_actions.append(f"MCP {operation}")

        elif tool_name == "TodoWrite":
            todos = tool_input.get("todos", [])
            key_actions.append(f"Created {len(todos)} todos")

        else:
            key_actions.append(f"Used {tool_name}")

    # Deduplicate key actions
    key_actions = (
        list(dict.fromkeys(key_actions))
        if key_actions
        else ["No specific actions detected"]
    )

    # Deduplicate file paths
    for op_type in file_operations:
        file_operations[op_type] = list(dict.fromkeys(file_operations[op_type]))

    # Extract result summary (first 200 chars)
    result_summary = (
        (session.result[:200] + "...") if len(session.result) > 200 else session.result
    )

    return {
        "tool_counts": tool_counts,
        "tool_details": tool_details,
        "file_operations": file_operations,
        "key_actions": key_actions,
        "total_tool_calls": sum(tool_counts.values()),
        "result_summary": result_summary,
        "usage_stats": {
            "total_cost_usd": session.total_cost_usd,
            "num_turns": session.num_turns,
            "duration_ms": session.duration_ms,
            "duration_api_ms": session.duration_api_ms,
            **session.usage,
        },
    }


# --------------------------------------------------------------------------- NDJSON stream


async def _ndjson_from_cli(request: ClaudeCodeRequest):
    """
    Yields each JSON object emitted by the *claude-code* CLI.

    • Robust against UTF-8 splits across chunks (incremental decoder).
    • Robust against braces inside strings (uses json.JSONDecoder.raw_decode)
    • Falls back to `json_repair.repair_json` when necessary.
    """
    from json_repair import repair_json

    workspace = request.cwd()
    workspace.mkdir(parents=True, exist_ok=True)

    proc = await asyncio.create_subprocess_exec(
        CLAUDE_CLI,
        *request.as_cmd_args(),
        cwd=str(workspace),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        start_new_session=True,  # isolate from parent's SIGINT
    )

    decoder = codecs.getincrementaldecoder("utf-8")()
    json_decoder = json.JSONDecoder()
    buffer: str = ""  # text buffer that may hold >1 JSON objects

    try:
        while True:
            chunk = await proc.stdout.read(4096)
            if not chunk:
                break

            # 1) decode *incrementally* so we never split multibyte chars
            buffer += decoder.decode(chunk)

            # 2) try to peel off as many complete JSON objs as possible
            while buffer:
                buffer = buffer.lstrip()  # remove leading spaces/newlines
                if not buffer:
                    break
                try:
                    obj, idx = json_decoder.raw_decode(buffer)
                    yield obj
                    buffer = buffer[idx:]  # keep remainder for next round
                except json.JSONDecodeError:
                    # incomplete → need more bytes
                    break

        # 3) flush any tail bytes in the incremental decoder
        buffer += decoder.decode(b"", final=True)
        buffer = buffer.strip()
        if buffer:
            try:
                obj, idx = json_decoder.raw_decode(buffer)
                yield obj
            except json.JSONDecodeError:
                try:
                    fixed = repair_json(buffer)
                    yield json.loads(fixed)
                    log.warning("Repaired malformed JSON fragment at stream end")
                except Exception:
                    log.error("Skipped unrecoverable JSON tail: %.120s…", buffer)

        # 4) propagate non-zero exit code
        if await proc.wait() != 0:
            err = (await proc.stderr.read()).decode().strip()
            raise RuntimeError(err or "CLI exited non-zero")

    finally:
        with contextlib.suppress(ProcessLookupError):
            proc.terminate()
        try:
            await asyncio.wait_for(proc.wait(), timeout=5.0)
        except asyncio.TimeoutError:
            with contextlib.suppress(ProcessLookupError):
                proc.kill()
            with contextlib.suppress(Exception):
                await proc.wait()


# --------------------------------------------------------------------------- SSE route
async def stream_cc_cli_events(request: ClaudeCodeRequest):
    if not CLAUDE_CLI:
        raise RuntimeError(
            "Claude CLI binary not found (npm i -g @anthropic-ai/claude-code)"
        )
    async with contextlib.aclosing(_ndjson_from_cli(request)) as stream:
        async for obj in stream:
            yield obj
    yield {"type": "done"}


print_readable = partial(as_readable, md=True, display_str=True)


def _pp_system(sys_obj: dict[str, Any], theme) -> None:
    txt = (
        f"◼️  **Claude Code Session**  \n"
        f"- id: `{sys_obj.get('session_id', '?')}`  \n"
        f"- model: `{sys_obj.get('model', '?')}`  \n"
        f"- tools: {', '.join(sys_obj.get('tools', [])[:8])}"
        + ("…" if len(sys_obj.get("tools", [])) > 8 else "")
    )
    print_readable(txt, border=False, theme=theme)


def _pp_thinking(thought: str, theme) -> None:
    text = f"""
    🧠 Thinking:
    {thought}
    """
    print_readable(text, border=True, theme=theme)


def _pp_assistant_text(text: str, theme) -> None:
    txt = f"""
    > 🗣️ Claude:
    {text}
    """
    print_readable(txt, theme=theme)


def _pp_tool_use(tu: dict[str, Any], theme) -> None:
    preview = shorten(str(tu["input"]).replace("\n", " "), 130)
    body = f"- 🔧 Tool Use — {tu['name']}({tu['id']}) - input: {preview}"
    print_readable(body, border=False, panel=False, theme=theme)


def _pp_tool_result(tr: dict[str, Any], theme) -> None:
    body_preview = shorten(str(tr["content"]).replace("\n", " "), 130)
    status = "ERR" if tr.get("is_error") else "OK"
    body = (
        f"- 📄 Tool Result({tr['tool_use_id']}) - {status}\n\n\tcontent: {body_preview}"
    )
    print_readable(body, border=False, panel=False, theme=theme)


def _pp_final(sess: ClaudeSession, theme) -> None:
    usage = sess.usage or {}
    cost_str = (
        f"${sess.total_cost_usd:.4f}" if sess.total_cost_usd is not None else "N/A"
    )
    txt = (
        f"### ✅ Session complete - {datetime.now(timezone.utc).isoformat(timespec='seconds')} UTC\n"
        f"**Result:**\n\n{sess.result or ''}\n\n"
        f"- cost: **{cost_str}**  \n"
        f"- turns: **{sess.num_turns}**  \n"
        f"- duration: **{sess.duration_ms} ms** (API {sess.duration_api_ms} ms)  \n"
        f"- tokens in/out: {usage.get('input_tokens', 0)}/{usage.get('output_tokens', 0)}"
    )
    print_readable(txt, theme=theme)


# --------------------------------------------------------------------------- internal utils


async def _maybe_await(func, *args, **kw):
    """Call func which may be sync or async."""
    res = func(*args, **kw) if func else None
    if inspect.iscoroutine(res):
        await res


# --------------------------------------------------------------------------- main parser
async def stream_claude_code_cli(  # noqa: C901
    request: ClaudeCodeRequest,
    session: ClaudeSession | None = None,
    *,
    on_system: Callable[[dict[str, Any]], None] | None = None,
    on_thinking: Callable[[str], None] | None = None,
    on_text: Callable[[str], None] | None = None,
    on_tool_use: Callable[[dict[str, Any]], None] | None = None,
    on_tool_result: Callable[[dict[str, Any]], None] | None = None,
    on_final: Callable[[ClaudeSession], None] | None = None,
) -> AsyncIterator[ClaudeChunk | dict | ClaudeSession]:
    """
    Consume the ND-JSON stream produced by ndjson_from_cli()
    and return a fully-populated ClaudeSession.

    If callbacks are omitted a default pretty-print is emitted.
    """
    if session is None:
        session = ClaudeSession()
    theme = request.cli_display_theme or "light"

    stream = stream_cc_cli_events(request)
    try:
        async for obj in stream:
            typ = obj.get("type", "unknown")
            chunk = ClaudeChunk(raw=obj, type=typ)
            session.chunks.append(chunk)

            # ------------------------ SYSTEM -----------------------------------
            if typ == "system":
                data = obj
                session.session_id = data.get("session_id", session.session_id)
                session.model = data.get("model", session.model)
                await _maybe_await(on_system, data)
                if request.verbose_output:
                    _pp_system(data, theme)
                yield data

            # ------------------------ ASSISTANT --------------------------------
            elif typ == "assistant":
                msg = obj["message"]
                session.messages.append(msg)

                for blk in msg.get("content", []):
                    btype = blk.get("type")
                    if btype == "thinking":
                        thought = blk.get("thinking", "").strip()
                        chunk.thinking = thought
                        session.thinking_log.append(thought)
                        await _maybe_await(on_thinking, thought)
                        if request.verbose_output:
                            _pp_thinking(thought, theme)

                    elif btype == "text":
                        text = blk.get("text", "")
                        chunk.text = text
                        await _maybe_await(on_text, text)
                        if request.verbose_output:
                            _pp_assistant_text(text, theme)

                    elif btype == "tool_use":
                        tu = {
                            "id": blk["id"],
                            "name": blk["name"],
                            "input": blk["input"],
                        }
                        chunk.tool_use = tu
                        session.tool_uses.append(tu)
                        await _maybe_await(on_tool_use, tu)
                        if request.verbose_output:
                            _pp_tool_use(tu, theme)

                    elif btype == "tool_result":
                        tr = {
                            "tool_use_id": blk["tool_use_id"],
                            "content": blk["content"],
                            "is_error": blk.get("is_error", False),
                        }
                        chunk.tool_result = tr
                        session.tool_results.append(tr)
                        await _maybe_await(on_tool_result, tr)
                        if request.verbose_output:
                            _pp_tool_result(tr, theme)
                yield chunk

            # ------------------------ USER (tool_result containers) ------------
            elif typ == "user":
                msg = obj["message"]
                session.messages.append(msg)
                for blk in msg.get("content", []):
                    if blk.get("type") == "tool_result":
                        tr = {
                            "tool_use_id": blk["tool_use_id"],
                            "content": blk["content"],
                            "is_error": blk.get("is_error", False),
                        }
                        chunk.tool_result = tr
                        session.tool_results.append(tr)
                        await _maybe_await(on_tool_result, tr)
                        if request.verbose_output:
                            _pp_tool_result(tr, theme)
                yield chunk

            # ------------------------ RESULT -----------------------------------
            elif typ == "result":
                session.result = obj.get("result", "").strip()
                session.usage = obj.get("usage", {})
                session.total_cost_usd = obj.get("total_cost_usd")
                session.num_turns = obj.get("num_turns")
                session.duration_ms = obj.get("duration_ms")
                session.duration_api_ms = obj.get("duration_api_ms")
                session.is_error = obj.get("is_error", False)

            # ------------------------ DONE -------------------------------------
            elif typ == "done":
                break
    finally:
        await stream.aclose()

    # final pretty print
    await _maybe_await(on_final, session)
    if request.verbose_output:
        _pp_final(session, theme)

    yield session
