"""
.agents/skills/lib_py/layout.py
Shared 2xK grid TUI layout engine for multi-agent workspaces.
"""
from dataclasses import dataclass
from typing import List, Dict, Optional, Any
import json
import sys
import os
import argparse


@dataclass
class PaneInfo:
    pane_id: str
    x: int
    y: int
    width: int
    height: int
    # NOTE: no `focused` field. The 2xK engine is deliberately geometry- and
    # structure-driven so that identical pane sets always yield identical
    # decisions. Focus is user-interaction state and would make the result
    # non-deterministic; herdr still reports it in the payload if ever needed.


@dataclass
class LayoutDecision:
    target_pane_id: str
    direction: str          # 'right' | 'down' | 'overflow'
    is_overflow: bool = False
    reason: str = ""


def extract_panes(data: Dict[str, Any]) -> List[PaneInfo]:
    """Extracts list of PaneInfo from herdr layout JSON payload.

    Accepts all three shapes herdr 0.8 emits: result.layout.panes,
    result.panes, and a bare top-level panes array.
    """
    if not isinstance(data, dict):
        return []

    res = data.get("result", {})
    if not isinstance(res, dict):
        res = {}

    layout = res.get("layout", {})
    if isinstance(layout, dict) and "panes" in layout:
        raw_panes = layout.get("panes", [])
    else:
        raw_panes = res.get("panes", []) or data.get("panes", [])

    panes: List[PaneInfo] = []
    for p in raw_panes:
        if not isinstance(p, dict):
            continue
        pid = str(p.get("pane_id", ""))
        rect = p.get("rect", {})
        if not isinstance(rect, dict):
            rect = {}
        x = int(rect.get("x", 0))
        y = int(rect.get("y", 0))
        w = int(rect.get("width", 0))
        h = int(rect.get("height", 0))
        if pid:
            panes.append(PaneInfo(pane_id=pid, x=x, y=y, width=w, height=h))

    return panes


def _pane_id(pane: Any) -> str:
    if isinstance(pane, PaneInfo):
        return pane.pane_id
    return str(pane or "")


def _right(reason: str, pane: Any) -> LayoutDecision:
    return LayoutDecision(target_pane_id=_pane_id(pane), direction="right", reason=reason)


def _down(reason: str, pane: Any) -> LayoutDecision:
    return LayoutDecision(target_pane_id=_pane_id(pane), direction="down", reason=reason)


def _overflow(reason: str, pane: Any) -> LayoutDecision:
    return LayoutDecision(
        target_pane_id=_pane_id(pane), direction="overflow", is_overflow=True, reason=reason
    )


def _full_height_pane(col: List[PaneInfo], area_h: int, tol: int = 2) -> Optional[PaneInfo]:
    """Single pane that occupies the full column height. Else None (R-2)."""
    if len(col) != 1:
        return None
    if area_h <= 0:
        return col[0]
    return col[0] if abs(col[0].height - area_h) <= tol else None


def _group_columns(panes: List[PaneInfo]) -> List[List[PaneInfo]]:
    columns: List[List[PaneInfo]] = []
    for p in sorted(panes, key=lambda q: (q.x, q.y)):
        matched = False
        for col in columns:
            if abs(col[0].x - p.x) <= 2:
                col.append(p)
                matched = True
                break
        if not matched:
            columns.append([p])
    for col in columns:
        col.sort(key=lambda q: q.y)
    return columns


def _area_height(data: Dict[str, Any], panes: List[PaneInfo]) -> int:
    res = data.get("result") if isinstance(data, dict) else {}
    if not isinstance(res, dict):
        res = {}
    layout = res.get("layout")
    if isinstance(layout, dict):
        area = layout.get("area")
        if isinstance(area, dict) and area.get("height"):
            try:
                return int(area["height"])
            except (TypeError, ValueError):
                pass
    if not panes:
        return 0
    return max(p.y + p.height for p in panes)


def _decide(
    cols: List[List[PaneInfo]],
    area_h: int,
    max_columns: int,
    max_rows: int,
    min_cols: int,
    min_rows: int,
) -> LayoutDecision:
    """Shared decision table (GUI observation)."""
    if not cols:
        return _right("no_panes_default", "")

    # ① New column only from a full-height pane (R-2).
    if len(cols) < max_columns:
        fh = _full_height_pane(cols[-1], area_h)
        if fh is not None:
            if fh.width > 0 and fh.width // 2 < min_cols:
                return _overflow("column_width_overflow", fh)
            return _right("new_column_right", fh)

    # ② Fill the shortest column (leftmost on ties).
    shortest = min(cols, key=lambda c: (len(c), c[0].x))
    if len(shortest) < max_rows:
        bottom = shortest[-1]
        if min_rows > 0 and bottom.height > 0 and bottom.height // 2 < min_rows:
            return _overflow("row_height_overflow", bottom)
        return _down("fill_column", bottom)

    # ③ Capacity reached.
    return _overflow("grid_capacity_reached", cols[-1][0])


def _decide_headless(
    panes: List[PaneInfo], max_columns: int, max_rows: int
) -> LayoutDecision:
    """Same table as _decide, observed from creation order (0x0 panes)."""
    n = len(panes)
    if n >= max_columns * max_rows:
        return _overflow("grid_capacity_reached", panes[-1])
    if n < max_columns:
        return _right("new_column_right", panes[n - 1])
    return _down("fill_column", panes[n - max_columns])


def compute_2xk_layout(
    data: Dict[str, Any],
    min_cols: int = 15,
    min_rows: int = 0,
    max_columns: Optional[int] = 2,
    max_rows: Optional[int] = 2,
    default_anchor_id: Optional[str] = None
) -> LayoutDecision:
    """Balanced 2xK grid. Herdr-supported splits only: 'right' and 'down'.

    GUI and headless share one decision table. Observation differs:
    geometry grouping vs creation-order column occupancy.
    """
    if max_columns is None:
        max_columns = 2
    if max_rows is None:
        max_rows = 2

    panes = extract_panes(data)

    if not panes:
        return _right("no_panes_default", default_anchor_id or "")

    if all(p.width <= 0 or p.height <= 0 for p in panes):
        return _decide_headless(panes, max_columns, max_rows)

    cols = _group_columns(panes)
    return _decide(cols, _area_height(data, panes), max_columns, max_rows, min_cols, min_rows)


def _env_int(*names: str, default: Optional[int] = None) -> Optional[int]:
    """First *valid* int among the env vars in *names*, else `default`.

    `default` is an explicit parameter rather than an `or` at the call site so a
    legitimate 0 survives (MAM_MIN_PANE_COLS=0 means 0, not the 15 default).

    An unparsable value is skipped rather than raised or treated as terminal: a
    typo in an operator's shell must not take the whole layout call down (lib.sh
    would silently fall back to 'right'), and must not shadow a later candidate
    that IS set correctly -- MAM_MIN_COLS is a legacy alias while
    MAM_MIN_PANE_COLS is the name .mam.env.example documents, so aborting on the
    first bad value would discard the documented setting. Empty values already
    fell through; this makes invalid values behave the same way.
    """
    for n in names:
        raw = os.environ.get(n, "").strip()
        if raw:
            try:
                return int(raw)
            except ValueError:
                continue
    return default


def main():
    parser = argparse.ArgumentParser(description="Compute 2xK grid TUI layout split direction")
    parser.add_argument("--min-cols", type=int, default=_env_int("MAM_MIN_COLS", "MAM_MIN_PANE_COLS", default=15))
    parser.add_argument("--min-rows", type=int, default=_env_int("MAM_MIN_ROWS", "MAM_MIN_PANE_ROWS", default=0))
    parser.add_argument("--max-cols", type=int, default=_env_int("MAM_MAX_COLS", "MAM_MAX_PANE_COLS", default=2))
    parser.add_argument("--max-rows", type=int, default=_env_int("MAM_MAX_ROWS", "MAM_MAX_PANE_ROWS", default=2))
    parser.add_argument("--sample-pane", type=str, default=None)
    parser.add_argument("--json", action="store_true", help="Output full JSON decision")

    args = parser.parse_args()

    raw_input = sys.stdin.read().strip()
    data = {}
    if raw_input:
        try:
            data = json.loads(raw_input)
        except Exception:
            data = {}

    decision = compute_2xk_layout(
        data=data,
        min_cols=args.min_cols,
        min_rows=args.min_rows,
        max_columns=args.max_cols,
        max_rows=args.max_rows,
        default_anchor_id=args.sample_pane
    )

    if args.json:
        print(json.dumps({
            "target_pane_id": decision.target_pane_id,
            "direction": decision.direction,
            "is_overflow": decision.is_overflow,
            "reason": decision.reason
        }))
    else:
        if decision.target_pane_id:
            print(f"{decision.direction} {decision.target_pane_id}")
        else:
            print(decision.direction)


if __name__ == "__main__":
    main()