"""Core run() function — streams a Claude Agent SDK session with retry and timeout."""

from __future__ import annotations

import asyncio
import random
import time
from collections.abc import Callable
from pathlib import Path

from claude_runner._types import RunResult


async def run(
    prompt: str,
    *,
    model: str | None = None,
    system_prompt: str | None = None,
    permission_mode: str | None = None,
    cwd: str | Path | None = None,
    max_turns: int | None = 1,
    tools: list[str] | None = [],
    thinking: dict | None = None,
    setting_sources: list[str] | None = None,
    retries: int = 0,
    timeout_minutes: float = 30,
    retry_base_delay: float = 10.0,
    retry_max_delay: float | None = None,
    retry_jitter: bool = False,
    on_text: Callable[[str], None] | None = None,
    on_stderr: Callable[[str], None] | None = None,
    options: object | None = None,
) -> RunResult:
    """Run a Claude session via the Agent SDK.

    Simple cases use flat kwargs; pass ``options`` (a ``ClaudeAgentOptions``)
    for full SDK control. Flat kwargs override matching fields on ``options``.

    Args:
        prompt: The user prompt to send.
        model: Model identifier (e.g. "claude-sonnet-4-6").
        system_prompt: System prompt prepended to the conversation.
        permission_mode: SDK permission mode (e.g. "bypassPermissions").
        cwd: Working directory for the Claude CLI subprocess.
        max_turns: Maximum conversational turns (default 1 for direct LLM call).
        tools: List of tool names to enable (default [] = no tools; pass None for SDK defaults).
        thinking: Thinking configuration dict (e.g. {"type": "disabled"}).
        setting_sources: List of setting sources (e.g. ["user", "project"]).
        retries: Number of *additional* attempts after the first (0 = no retry).
        timeout_minutes: Per-attempt timeout in minutes.
        retry_base_delay: Base delay for exponential backoff (base * 2^attempt).
        retry_max_delay: Cap on the computed delay (None = no cap).
        retry_jitter: Multiply delay by random(0.75, 1.25) to avoid thundering herd.
        on_text: Callback for each TextBlock received.
        on_stderr: Callback for each stderr line from the CLI subprocess.
        options: Full ClaudeAgentOptions escape hatch. Flat kwargs override.

    Returns:
        RunResult with session output, cost, and metadata. Never raises for
        session failures — check ``is_error`` instead. Only raises for
        programmer errors (e.g. SDK not installed).
    """
    from claude_agent_sdk import (
        AssistantMessage,
        ClaudeAgentOptions,
        ResultMessage,
        TextBlock,
        query,
    )

    opts = _build_options(
        ClaudeAgentOptions,
        base=options,
        model=model,
        system_prompt=system_prompt,
        permission_mode=permission_mode,
        cwd=str(cwd) if cwd is not None else None,
        max_turns=max_turns,
        tools=tools,
        thinking=thinking,
        setting_sources=setting_sources,
        on_stderr=on_stderr,
    )

    total_attempts = 1 + retries
    last_error: Exception | None = None

    for attempt in range(total_attempts):
        t0 = time.monotonic()
        try:
            result = await asyncio.wait_for(
                _run_single_attempt(
                    query, opts, prompt,
                    AssistantMessage=AssistantMessage,
                    ResultMessage=ResultMessage,
                    TextBlock=TextBlock,
                    on_text=on_text,
                ),
                timeout=timeout_minutes * 60,
            )
            result.duration_ms = int((time.monotonic() - t0) * 1000)
            return result

        except asyncio.TimeoutError:
            last_error = TimeoutError(f"Session timed out after {timeout_minutes}m")

        except Exception as exc:
            last_error = exc

        if attempt < total_attempts - 1:
            delay = retry_base_delay * (2 ** attempt)
            if retry_max_delay is not None:
                delay = min(delay, retry_max_delay)
            if retry_jitter:
                delay *= random.uniform(0.75, 1.25)
            await asyncio.sleep(delay)

    # All attempts exhausted.
    duration_ms = int((time.monotonic() - t0) * 1000)  # type: ignore[possibly-undefined]
    return RunResult(
        text="",
        is_error=True,
        error=f"Failed after {total_attempts} attempt(s): {last_error}",
        duration_ms=duration_ms,
    )


def _build_options(
    ClaudeAgentOptions: type,
    *,
    base: object | None,
    model: str | None,
    system_prompt: str | None,
    permission_mode: str | None,
    cwd: str | None,
    max_turns: int | None,
    tools: list[str] | None,
    thinking: dict | None,
    setting_sources: list[str] | None,
    on_stderr: Callable[[str], None] | None,
) -> object:
    """Merge flat kwargs onto a base ClaudeAgentOptions."""
    if base is not None:
        opts = base
    else:
        opts = ClaudeAgentOptions()

    # Apply flat overrides.
    if model is not None:
        opts.model = model  # type: ignore[attr-defined]
    if system_prompt is not None:
        opts.system_prompt = system_prompt  # type: ignore[attr-defined]
    if permission_mode is not None:
        opts.permission_mode = permission_mode  # type: ignore[attr-defined]
    if cwd is not None:
        opts.cwd = cwd  # type: ignore[attr-defined]
    if max_turns is not None:
        opts.max_turns = max_turns  # type: ignore[attr-defined]
    if tools is not None:
        opts.tools = tools  # type: ignore[attr-defined]
    if thinking is not None:
        opts.thinking = thinking  # type: ignore[attr-defined]
    if setting_sources is not None:
        opts.setting_sources = setting_sources  # type: ignore[attr-defined]
    if on_stderr is not None:
        opts.stderr = on_stderr  # type: ignore[attr-defined]

    return opts


async def _run_single_attempt(
    query_fn: Callable,
    opts: object,
    prompt: str,
    *,
    AssistantMessage: type,
    ResultMessage: type,
    TextBlock: type,
    on_text: Callable[[str], None] | None,
) -> RunResult:
    """Execute one SDK query attempt. Returns RunResult or raises."""
    collected: list[str] = []
    all_messages: list = []
    result_msg = None

    stream = query_fn(prompt=prompt, options=opts)

    # Manual async iteration so we can catch/skip unknown message types
    # without aborting the whole stream.
    aiter = stream.__aiter__()
    while True:
        try:
            message = await aiter.__anext__()
        except StopAsyncIteration:
            break
        except StopIteration:
            break
        except Exception as exc:
            if "Unknown message type" in str(exc):
                continue
            raise

        all_messages.append(message)

        if isinstance(message, AssistantMessage):
            # SDK sends AssistantMessage with error="rate_limit" (empty content)
            # when the API returns 429. Raise to trigger retry.
            if getattr(message, "error", None) == "rate_limit":
                raise RuntimeError("API rate limited (AssistantMessage.error)")

            for block in message.content:
                if isinstance(block, TextBlock):
                    if on_text is not None:
                        on_text(block.text)
                    collected.append(block.text)

        elif isinstance(message, ResultMessage):
            result_msg = message
            if message.result and not collected:
                if on_text is not None:
                    on_text(message.result)
                collected.append(message.result)

    # Stream ended — build result.
    if not collected and result_msg is None:
        raise RuntimeError("Stream ended with no content (no ResultMessage)")
    if not collected and result_msg is not None and result_msg.is_error:
        raise RuntimeError(f"Stream ended with error: {result_msg.result}")

    text = "\n".join(collected)

    # Extract fields from ResultMessage when available.
    cost_usd = 0.0
    usage: dict = {}
    num_turns = 0
    session_id = ""
    is_error = False
    error: str | None = None

    if result_msg is not None:
        cost_usd = getattr(result_msg, "total_cost_usd", 0.0) or 0.0
        usage = getattr(result_msg, "usage", {}) or {}
        num_turns = getattr(result_msg, "num_turns", 0) or 0
        session_id = getattr(result_msg, "session_id", "") or ""
        is_error = getattr(result_msg, "is_error", False) or False
        if is_error:
            error = getattr(result_msg, "result", None) or "Unknown SDK error"

    return RunResult(
        text=text,
        cost_usd=cost_usd,
        usage=usage,
        num_turns=num_turns,
        session_id=session_id,
        is_error=is_error,
        error=error,
        result_message=result_msg,
        messages=all_messages,
    )


def run_sync(prompt: str, **kwargs) -> RunResult:
    """Synchronous wrapper around run(). Cannot be called from a running event loop."""
    return asyncio.run(run(prompt, **kwargs))
