from __future__ import annotations

import io
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel

from ...items import ItemHelpers, RunItem, ToolApprovalItem, TResponseInputItem
from ...result import RunResultBase, RunResultStreaming
from ...run_internal.items import run_items_to_input_items
from ...util._json import _to_dump_compatible
from ..errors import WorkspaceReadNotFoundError
from ..session.base_sandbox_session import BaseSandboxSession

_EXCLUDED_MEMORY_ITEM_TYPES = frozenset(
    {
        "compaction",
        "image_generation_call",
        "reasoning",
    }
)
_INCLUDED_MEMORY_ITEM_TYPES = frozenset(
    {
        "apply_patch_call",
        "apply_patch_call_output",
        "computer_call",
        "computer_call_output",
        "custom_tool_call",
        "custom_tool_call_output",
        "function_call",
        "function_call_output",
        "local_shell_call",
        "local_shell_call_output",
        "mcp_approval_request",
        "mcp_approval_response",
        "mcp_call",
        "shell_call",
        "shell_call_output",
        "tool_search_call",
        "tool_search_output",
        "web_search_call",
    }
)


def _validate_relative_path(*, name: str, path: Path) -> None:
    if path.is_absolute():
        raise ValueError(f"{name} must be relative to the sandbox workspace root, got: {path}")
    if ".." in path.parts:
        raise ValueError(f"{name} must not escape root, got: {path}")
    if path.parts in [(), (".",)]:
        raise ValueError(f"{name} must be non-empty")


class RolloutTerminalMetadata(BaseModel):
    terminal_state: Literal[
        "completed",
        "interrupted",
        "cancelled",
        "failed",
        "max_turns_exceeded",
        "guardrail_tripped",
    ]
    exception_type: str | None = None
    exception_message: str | None = None
    has_final_output: bool = False


def dump_rollout_json(result: Any) -> str:
    return json.dumps(result, separators=(",", ":")) + "\n"


def _normalize_jsonl_line(*, rollout_contents: str) -> bytes:
    try:
        obj = json.loads(rollout_contents)
    except Exception as exc:
        raise ValueError("rollout_contents must be valid JSON text") from exc
    line = json.dumps(obj, separators=(",", ":"))
    return (line + "\n").encode("utf-8")


def _should_include_memory_item(item: TResponseInputItem) -> bool:
    role = item.get("role")
    if role in {"developer", "system"}:
        return False
    if role in {"assistant", "tool", "user"}:
        return True

    item_type = item.get("type")
    if item_type in _EXCLUDED_MEMORY_ITEM_TYPES:
        return False
    return item_type in _INCLUDED_MEMORY_ITEM_TYPES


def _sanitize_memory_items(items: list[TResponseInputItem]) -> list[TResponseInputItem]:
    return [item for item in items if _should_include_memory_item(item)]


async def write_rollout(
    *,
    session: BaseSandboxSession,
    rollout_contents: str,
    rollouts_path: str = "sessions",
    file_name: str | None = None,
) -> Path:
    rollouts_dir_rel = Path(rollouts_path)
    _validate_relative_path(name="rollouts_path", path=rollouts_dir_rel)
    line_bytes = _normalize_jsonl_line(rollout_contents=rollout_contents)

    if file_name is not None:
        requested_file_rel = Path(file_name.strip())
        if not requested_file_rel.name.endswith(".jsonl") or len(requested_file_rel.parts) != 1:
            raise ValueError("file_name must be a simple .jsonl filename")
        dest_file_path_rel = rollouts_dir_rel / requested_file_rel
    else:
        dest_file_path_rel = None
        for _ in range(10):
            rollout_id = str(uuid.uuid4())
            candidate_rel = rollouts_dir_rel / f"{rollout_id}.jsonl"
            prior_bytes = await _read_existing_bytes(session=session, path=candidate_rel)
            if prior_bytes is None:
                dest_file_path_rel = candidate_rel
                break
    if dest_file_path_rel is None:
        raise ValueError(f"failed to allocate a unique rollout id under: {rollouts_dir_rel}")

    await session.mkdir(dest_file_path_rel.parent, parents=True)
    prior_bytes = await _read_existing_bytes(session=session, path=dest_file_path_rel)
    if prior_bytes is None:
        await session.write(dest_file_path_rel, io.BytesIO(line_bytes))
    else:
        await session.write(dest_file_path_rel, io.BytesIO(prior_bytes + line_bytes))
    return dest_file_path_rel


async def _read_existing_bytes(*, session: BaseSandboxSession, path: Path) -> bytes | None:
    try:
        handle = await session.read(path)
    except WorkspaceReadNotFoundError:
        return None

    try:
        payload = handle.read()
    finally:
        handle.close()
    return payload.encode("utf-8") if isinstance(payload, str) else bytes(payload)


def terminal_metadata_for_result(
    result: RunResultBase,
    *,
    exception: BaseException | None = None,
) -> RolloutTerminalMetadata:
    if result.final_output is not None:
        return RolloutTerminalMetadata(terminal_state="completed", has_final_output=True)
    if getattr(result, "interruptions", None):
        return RolloutTerminalMetadata(terminal_state="interrupted", has_final_output=False)

    exc = exception
    if exc is None and isinstance(result, RunResultStreaming):
        exc = getattr(result, "_stored_exception", None)
        if exc is None and result._cancel_mode == "immediate":
            return RolloutTerminalMetadata(terminal_state="cancelled", has_final_output=False)

    if exc is None:
        return RolloutTerminalMetadata(terminal_state="failed", has_final_output=False)

    return terminal_metadata_for_exception(exc)


def terminal_metadata_for_exception(exc: BaseException) -> RolloutTerminalMetadata:
    exc_name = type(exc).__name__
    terminal_state: Literal[
        "max_turns_exceeded",
        "guardrail_tripped",
        "cancelled",
        "failed",
    ]
    if exc_name == "MaxTurnsExceeded":
        terminal_state = "max_turns_exceeded"
    elif "Guardrail" in exc_name:
        terminal_state = "guardrail_tripped"
    elif exc_name == "CancelledError":
        terminal_state = "cancelled"
    else:
        terminal_state = "failed"
    return RolloutTerminalMetadata(
        terminal_state=terminal_state,
        exception_type=exc_name,
        exception_message=str(exc) or None,
        has_final_output=False,
    )


def _serialize_interruption_raw_item(raw_item: Any) -> Any:
    if isinstance(raw_item, BaseModel):
        return _to_dump_compatible(raw_item.model_dump(exclude_unset=True))
    if isinstance(raw_item, dict):
        return dict(raw_item)
    return _to_dump_compatible(raw_item)


def build_rollout_payload(
    *,
    input: str | list[TResponseInputItem],
    new_items: list[RunItem],
    final_output: Any,
    interruptions: list[ToolApprovalItem],
    terminal_metadata: RolloutTerminalMetadata,
) -> dict[str, Any]:
    input_items = _sanitize_memory_items(ItemHelpers.input_to_new_input_list(input))
    generated_items = _to_dump_compatible(
        _sanitize_memory_items(run_items_to_input_items(new_items))
    )

    serialized_interruptions = [
        _serialize_interruption_raw_item(interruption.raw_item) for interruption in interruptions
    ]

    payload: dict[str, Any] = {
        "updated_at": datetime.now(tz=timezone.utc).isoformat(),
        "input": _to_dump_compatible(input_items),
        "generated_items": generated_items,
    }
    if serialized_interruptions:
        payload["interruptions"] = serialized_interruptions
    payload["terminal_metadata"] = terminal_metadata.model_dump(mode="json")
    if final_output is not None:
        payload["final_output"] = _to_dump_compatible(final_output)
    return payload


def build_rollout_payload_from_result(
    result: RunResultBase,
    *,
    exception: BaseException | None = None,
    input_override: str | list[TResponseInputItem] | None = None,
) -> dict[str, Any]:
    interruptions = list(getattr(result, "interruptions", []))
    return build_rollout_payload(
        input=input_override if input_override is not None else result.input,
        new_items=result.new_items,
        final_output=result.final_output,
        interruptions=interruptions,
        terminal_metadata=terminal_metadata_for_result(result, exception=exception),
    )
