from __future__ import annotations

from typing import Any


def _is_known_time(value: Any) -> bool:
    return bool(value) and str(value).strip().lower() not in {"unknown", "none", "nan", "nat"}


def _stage_order(row: dict[str, Any], fallback_index: int = 0) -> int:
    value = row.get("stage_index", row.get("stage_order", fallback_index))
    try:
        return int(value)
    except (TypeError, ValueError):
        return fallback_index


def _stage_label(row: dict[str, Any]) -> str:
    for field in ("stage_title", "stage_label_public", "stage_label", "stage_id"):
        value = str(row.get(field) or "").strip()
        if value:
            return value
    return "Unnamed stage"


def prepare_gantt_rows(stage_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    ordered = sorted(
        stage_rows,
        key=lambda row: (_stage_order(row), str(row.get("stage_id", ""))),
    )
    calendar_axis = all(
        _is_known_time(row.get("stage_start_time")) and _is_known_time(row.get("stage_end_time"))
        for row in ordered
    )
    prepared: list[dict[str, Any]] = []
    for fallback_index, row in enumerate(ordered, start=1):
        start = row.get("stage_start_time")
        end = row.get("stage_end_time")
        stage_order = _stage_order(row, fallback_index)
        if calendar_axis:
            display_start = start
            display_end = end
            axis_mode = "calendar"
        else:
            display_start = stage_order
            display_end = display_start + 0.85
            axis_mode = "relative_order"

        time_note = row.get("temporal_anchor_summary") or ""
        if not time_note and int(row.get("known_action_time_anchor_count") or 0) > 0:
            time_note = "Action-level time anchors available"

        prepared.append(
            {
                **row,
                "stage_label": _stage_label(row),
                "stage_order": stage_order,
                "display_start": display_start,
                "display_end": display_end,
                "axis_mode": axis_mode,
                "time_note": time_note,
            }
        )
    return prepared


def build_timeline_figure(stage_rows: list[dict[str, Any]], event_id: str):
    import pandas as pd
    import plotly.express as px

    prepared = prepare_gantt_rows(stage_rows)
    if not prepared:
        return None

    frame = pd.DataFrame(prepared)
    if set(frame["axis_mode"]) == {"calendar"}:
        fig = px.timeline(
            frame,
            x_start="display_start",
            x_end="display_end",
            y="stage_label",
            color="stage_label",
            hover_data=["stage_id", "stage_order", "time_note"],
            title=f"{event_id}: public stage timeline",
        )
        fig.update_yaxes(autorange="reversed")
        return fig

    fig = px.bar(
        frame,
        x=[row["display_end"] - row["display_start"] for row in prepared],
        y="stage_label",
        base="display_start",
        orientation="h",
        color="stage_label",
        hover_data=["stage_id", "stage_order", "time_note"],
        title=f"{event_id}: relative stage order",
    )
    fig.update_yaxes(autorange="reversed")
    fig.update_layout(xaxis_title="Relative stage order")
    return fig