"""Project the latest run-scoped progress snapshot from append-only messages."""
from __future__ import annotations

import json
from collections.abc import Iterable
from typing import Any
from uuid import UUID

from sqlalchemy import select

from core.storage.models import Message

_VALID_STATUSES = {"pending", "in_progress", "completed"}


def _normalize_step(value: Any) -> dict[str, str] | None:
    if not isinstance(value, dict):
        return None
    step_id = str(value.get("id") or "").strip()
    title = str(value.get("title") or "").strip()
    status = str(value.get("status") or "pending")
    if not step_id or not title:
        return None
    if status not in _VALID_STATUSES:
        status = "pending"
    return {"id": step_id, "title": title, "status": status}


def _heal_monotonic(steps: list[dict[str, str]]) -> list[dict[str, str]]:
    last_completed = max(
        (i for i, step in enumerate(steps) if step["status"] == "completed"),
        default=-1,
    )
    return [
        {**step, "status": "completed"} if i < last_completed else dict(step)
        for i, step in enumerate(steps)
    ]


def apply_progress_args(
    current: list[dict[str, str]], args: Any,
) -> list[dict[str, str]]:
    """Apply current full snapshots plus legacy set/update/clear calls."""
    if not isinstance(args, dict):
        return [dict(step) for step in current]
    action = args.get("action") or ""
    if action == "clear":
        return []
    if isinstance(args.get("steps"), list):
        normalized: list[dict[str, str]] = []
        for raw in args["steps"]:
            step = _normalize_step(raw)
            if step is not None:
                normalized.append(step)
        return _heal_monotonic(normalized)
    if action != "update_step" or not isinstance(args.get("step"), dict):
        return [dict(step) for step in current]

    raw = args["step"]
    step_id = str(raw.get("id") or "").strip()
    if not step_id:
        return [dict(step) for step in current]
    next_steps: list[dict[str, str]] = []
    found = False
    for step in current:
        if step["id"] != step_id:
            next_steps.append(dict(step))
            continue
        found = True
        status = str(raw.get("status") or step["status"])
        next_steps.append({
            "id": step_id,
            "title": str(raw.get("title") or step["title"]).strip(),
            "status": status if status in _VALID_STATUSES else "pending",
        })
    if not found:
        normalized = _normalize_step(raw)
        if normalized is not None:
            next_steps.append(normalized)
    return _heal_monotonic(next_steps)


def project_progress_payloads(
    payloads: Iterable[dict[str, Any]],
) -> tuple[list[dict[str, str]], bool]:
    steps: list[dict[str, str]] = []
    seen = False
    for payload in payloads:
        if not isinstance(payload, dict) or payload.get("role") != "assistant":
            continue
        for call in payload.get("tool_calls") or []:
            function = call.get("function") if isinstance(call, dict) else None
            if not isinstance(function, dict) or function.get("name") != "task_progress":
                continue
            raw_args = function.get("arguments") or "{}"
            try:
                args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
            except (TypeError, json.JSONDecodeError):
                args = {}
            steps = apply_progress_args(steps, args)
            seen = True
    return steps, seen


def progress_waiting_for_user(payloads: Iterable[dict[str, Any]]) -> bool:
    for payload in payloads:
        if not isinstance(payload, dict) or payload.get("role") != "assistant":
            continue
        for call in payload.get("tool_calls") or []:
            function = call.get("function") if isinstance(call, dict) else None
            if isinstance(function, dict) and function.get("name") == "ask_user":
                return True
    return False


def latest_run_progress(session, task_id: UUID) -> dict[str, Any] | None:
    """Return the latest user message id and that run's projected plan."""
    run = session.execute(
        select(Message.message_id, Message.idx)
        .where(
            Message.task_id == task_id,
            Message.payload["role"].astext == "user",
        )
        .order_by(Message.idx.desc())
        .limit(1)
    ).first()
    if run is None:
        return None
    payloads = session.execute(
        select(Message.payload)
        .where(Message.task_id == task_id, Message.idx > run.idx)
        .order_by(Message.idx)
    ).scalars().all()
    steps, seen = project_progress_payloads(payloads)
    if not seen:
        return None
    return {
        "run_id": str(run.message_id),
        "steps": steps,
        "waiting": progress_waiting_for_user(payloads),
    }