"""Shared layout block implementation."""

from __future__ import annotations

from collections.abc import Callable
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from typing import Any

from yera.events.blocks.base.base import (
    _StreamHandle,
    create_event,
    get_parent_block_id,
)
from yera.events.models import (
    ColumnsData,
    ContainerData,
    LayoutEndData,
)
from yera.events.stream import await_input, push_output


@dataclass
class LayoutContext:
    """Context for tracking active layout blocks."""

    block_id: str
    block_type: str  # "container", "columns", or "form"
    column_index: int = 0  # For columns, tracks current column index
    input_count: int = (
        0  # Number of await_user blocks emitted in this layout; drained on exit
    )
    result_parsers: list[Callable[[str], Any]] = field(default_factory=list)


# Context variable to track the current layout
_current_layout: ContextVar[LayoutContext | None] = ContextVar(
    "current_layout", default=None
)


class LayoutHandle(_StreamHandle):
    """Stream handle for layout blocks (container, columns)."""

    def __init__(
        self,
        block_id: str,
        block_type: str,
        data: ContainerData | ColumnsData,
    ):
        super().__init__(block_id, block_type)
        self.data = data
        self._token: Token[LayoutContext | None] | None = None

    def __enter__(self):
        # Get parent block ID before setting new layout context
        parent_block_id = get_parent_block_id()

        # Create layout context (keep reference so __exit__ can drain input_count)
        layout_context = LayoutContext(
            block_id=self.block_id, block_type=self.block_type
        )
        self._layout_context = layout_context
        self._token = _current_layout.set(layout_context)

        # Emit layout event with pre-built data from factory
        event = self._create_chunk_event(
            data=self.data,
            message_type="layout",
            parent_block_id=parent_block_id,
        )
        push_output(event)

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # Clear layout context
        if self._token is not None:
            _current_layout.reset(self._token)
        # Emit layout_end so the stream can pause after all layout children have been sent
        layout_end_event = create_event(
            block_type="layout_end",
            block_id=self.block_id,
            data=LayoutEndData(),
            message_type="informational",
            chunk_id=1,
        )
        push_output(layout_end_event)
        # For container/columns, drain inputs so the agent does not block. Form drains in result().
        if self.block_type != "form":
            for _ in range(self._layout_context.input_count):
                await_input()
        return False

    def result(self) -> tuple[Any, ...]:
        """Return parsed form values. Only supported for form blocks."""
        if self.block_type != "form":
            raise TypeError("result() is only supported for form blocks")
        n = self._layout_context.input_count
        parsers = self._layout_context.result_parsers
        out: list[Any] = []
        for i in range(n):
            ev = await_input()
            if i < len(parsers):
                out.append(parsers[i](ev.data))
            else:
                out.append(ev.data)
        return tuple(out)


__all__ = ["LayoutContext", "LayoutHandle", "_current_layout"]
