"""Grading workbench screen."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from textual import on
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import Footer, Header, Input, Markdown, OptionList, Static
from textual.widgets.option_list import Option



import json as _json


def _format_content(value: Any) -> str:
    """Format trace input/output for markdown display."""
    if value is None:
        return "*—*"
    if isinstance(value, str):
        return value
    if isinstance(value, (dict, list)):
        return f"```json\n{_json.dumps(value, indent=2, ensure_ascii=False)}\n```"
    return str(value)


_NORMAL = "normal"
_SPAN_SELECT = "span_select"
_COMMENT = "comment"
_ADD_DATASET = "add_dataset"

_FILTER_NEEDS_REVIEW = "needs review"
_FILTER_JUDGE_ONLY = "judge only"
_FILTER_REVIEWED = "reviewed"
_FILTER_ALL = "all"
_FILTER_CYCLE = [_FILTER_NEEDS_REVIEW, _FILTER_JUDGE_ONLY, _FILTER_REVIEWED, _FILTER_ALL]

KIND_COLORS = {
    "llm": "cyan",
    "tool": "yellow",
    "retrieval": "green",
    "agent": "magenta",
    "embedding": "blue",
}


class SpanSelector(Static, can_focus=True):
    """Focusable widget for span tree — captures arrow/space keys in selection mode."""
    pass


class GradingScreen(Screen):
    BINDINGS = [
        Binding("y", "grade_good", "Good", show=True),
        Binding("n", "grade_bad", "Bad", show=True),
        Binding("s", "skip", "Skip", show=True),
        Binding("t", "toggle_filter", "Filter", show=True),
        Binding("c", "comment", "Comment", show=True),
        Binding("a", "add_to_dataset", "Add", show=True),
        Binding("d", "toggle_detail", "Detail", show=True),
        Binding("r", "refresh", "Refresh", show=True),
        Binding("b", "go_back", "Back", show=True),
        Binding("q", "quit", "Quit", show=True),
    ]

    show_detail: reactive[bool] = reactive(False)
    mode: reactive[str] = reactive(_NORMAL)
    span_cursor: reactive[int] = reactive(0)

    def __init__(self) -> None:
        super().__init__()
        self._traces: list[dict[str, Any]] = []
        self._current_idx: int = 0
        self._selected_spans: set[int] = set()
        self._comment: str | None = None
        self._picking_project: bool = False
        self._pre_comment_mode: str = _NORMAL
        self._filter: str = _FILTER_NEEDS_REVIEW
        self._human_graded_ids: set[str] = set()
        self._judge_graded_ids: set[str] = set()

    def compose(self) -> ComposeResult:
        yield Header()
        with Horizontal(id="grading-layout"):
            with Vertical(id="trace-sidebar"):
                yield Static("", id="grading-title")
                yield OptionList(id="trace-list")
            with VerticalScroll(id="detail-panel"):
                yield Markdown("", id="trace-input")
                yield Markdown("", id="trace-output")
                yield Static("", id="trace-summary")
                yield SpanSelector("", id="span-tree")
                yield Static("", id="comment-label")
                yield Input(
                    placeholder="Add a comment… (enter to save, esc to cancel)",
                    id="comment-input",
                )
        yield Footer()

    def on_mount(self) -> None:
        from ..app import BanditoApp

        self.query_one("#comment-input", Input).display = False
        self.query_one("#comment-label", Static).display = False

        app: BanditoApp = self.app
        projects = app.store.list_projects()

        if not projects:
            self.query_one("#grading-title", Static).update(
                "[red]No projects found[/red]"
            )
            return

        if app.project and app.project in projects:
            self._load_grading(app.project)
        elif len(projects) == 1:
            app.project = projects[0]
            self._load_grading(projects[0])
        else:
            self.query_one("#grading-title", Static).update(
                "[bold]Pick a project[/bold]"
            )
            option_list = self.query_one("#trace-list", OptionList)
            from ...retrieval import get_retrieval
            for p in projects:
                human_ids, _ = self._load_evals_from_client(p)
                r = get_retrieval(p)
                summaries = r.traces(tier="summary")
                ungraded = sum(1 for t in summaries if t["trace_id"] not in human_ids)
                option_list.add_option(
                    Option(f"{p}  [yellow]{ungraded} ungraded[/yellow]", id=p)
                )
            self._picking_project = True

    # ── Data Loading ───────────────────────────────────────────────────────

    @on(OptionList.OptionSelected, "#trace-list")
    def on_option_selected(self, event: OptionList.OptionSelected) -> None:
        from ..app import BanditoApp

        if self._picking_project and event.option and event.option.id:
            self._picking_project = False
            project = str(event.option.id)
            app: BanditoApp = self.app
            app.project = project
            self._load_grading(project)

    @on(OptionList.OptionHighlighted, "#trace-list")
    def on_trace_highlighted(self, event: OptionList.OptionHighlighted) -> None:
        if self._picking_project:
            return
        if event.option and event.option.id is not None:
            idx = int(str(event.option.id))
            self._show_trace(idx)

    def _load_evals_from_client(self, project: str) -> tuple[set[str], set[str]]:
        """Get trace IDs with human/judge grades. Returns (human_ids, judge_ids)."""
        from ...retrieval import get_retrieval

        r = get_retrieval(project)
        evals = r.evals()
        human_ids: set[str] = set()
        judge_ids: set[str] = set()
        for tid, ev in evals.items():
            if ev.get("human_score") is not None:
                human_ids.add(tid)
            if ev.get("judge_score") is not None:
                judge_ids.add(tid)
        return human_ids, judge_ids

    def _load_grading(self, project: str) -> None:
        from ...retrieval import get_retrieval

        r = get_retrieval(project)
        human_ids, judge_ids = self._load_evals_from_client(project)
        all_traces = r.traces(tier="trace")

        if self._filter == _FILTER_NEEDS_REVIEW:
            self._traces = [
                t for t in all_traces
                if t["trace_id"] not in human_ids
                and t["trace_id"] not in judge_ids
            ]
        elif self._filter == _FILTER_JUDGE_ONLY:
            self._traces = [
                t for t in all_traces
                if t["trace_id"] in judge_ids
                and t["trace_id"] not in human_ids
            ]
        elif self._filter == _FILTER_REVIEWED:
            self._traces = [
                t for t in all_traces
                if t["trace_id"] in human_ids
            ]
        else:
            self._traces = all_traces

        self._human_graded_ids = human_ids
        self._judge_graded_ids = judge_ids
        self._update_title()
        self._current_idx = 0
        self._populate_trace_list()
        if self._traces:
            self._show_trace(0)
        else:
            self.query_one("#trace-input", Markdown).update(
                "**All caught up!**"
                if self._filter in (_FILTER_NEEDS_REVIEW, _FILTER_JUDGE_ONLY)
                else "*No traces*"
            )
            self.query_one("#trace-output", Markdown).update("")
            self.query_one("#trace-summary", Static).update("")
            self.query_one("#span-tree", SpanSelector).update("")

    def _update_title(self) -> None:
        from ..app import BanditoApp

        app: BanditoApp = self.app
        labels = []
        for f in _FILTER_CYCLE:
            if f == self._filter:
                labels.append(f"[bold reverse] {f} [/bold reverse]")
            else:
                labels.append(f"[dim]{f}[/dim]")
        filter_bar = " · ".join(labels)
        self.query_one("#grading-title", Static).update(
            f"[bold]{app.project}[/bold] · {len(self._traces)} traces\n{filter_bar}"
        )

    def _populate_trace_list(self) -> None:
        option_list = self.query_one("#trace-list", OptionList)
        option_list.clear_options()
        for i, trace in enumerate(self._traces):
            trace_id = trace.get("trace_id", "")[:8]
            raw_input = trace.get("input", "")
            if isinstance(raw_input, str):
                input_text = raw_input.replace("\n", " ")[:24]
            elif isinstance(raw_input, (dict, list)):
                input_text = _json.dumps(raw_input, ensure_ascii=False)[:24]
            elif raw_input:
                input_text = str(raw_input)[:24]
            else:
                input_text = ""
            tid = trace.get("trace_id", "")
            has_human = tid in self._human_graded_ids
            has_judge = tid in self._judge_graded_ids
            if has_human:
                marker = "[green]✓[/green]"
            elif has_judge:
                marker = "[magenta]⚖[/magenta]"
            else:
                marker = " "
            label = f"{marker} {trace_id}  {input_text}"
            option_list.add_option(Option(label, id=str(i)))
        if self._traces:
            option_list.highlighted = 0

    # ── Trace Display ──────────────────────────────────────────────────────

    def _show_trace(self, idx: int) -> None:
        if idx < 0 or idx >= len(self._traces):
            return
        self._current_idx = idx
        trace = self._traces[idx]

        self.query_one("#trace-input", Markdown).update(
            f"### Input\n\n{_format_content(trace.get('input'))}"
        )
        self.query_one("#trace-output", Markdown).update(
            f"### Output\n\n{_format_content(trace.get('output'))}"
        )

        spans = trace.get("spans", [])
        total_tokens = sum(
            (s.get("input_tokens") or 0) + (s.get("output_tokens") or 0)
            for s in spans
        )
        total_cost = sum(s.get("cost") or 0 for s in spans)
        latency = self._calc_latency(trace)
        parts = [f"{len(spans)} spans"]
        if total_cost:
            parts.append(f"${total_cost:.4f}")
        if latency:
            parts.append(f"{latency:.1f}s")
        if total_tokens:
            parts.append(f"{total_tokens:,} tokens")
        self.query_one("#trace-summary", Static).update(
            f"[dim]{'  ·  '.join(parts)}[/dim]"
        )

        self._selected_spans.clear()
        self.span_cursor = 0
        self._render_span_tree()

    @staticmethod
    def _calc_latency(trace: dict[str, Any]) -> float | None:
        started = trace.get("started_at", "")
        ended = trace.get("ended_at", "")
        if started and ended:
            try:
                s = datetime.fromisoformat(started.replace("Z", "+00:00"))
                e = datetime.fromisoformat(ended.replace("Z", "+00:00"))
                return (e - s).total_seconds()
            except (ValueError, TypeError):
                pass
        return None

    # ── Span Tree Rendering ────────────────────────────────────────────────

    def _render_span_tree(self) -> None:
        span_widget = self.query_one("#span-tree", SpanSelector)
        if not self.show_detail:
            span_widget.update("")
            return

        if self._current_idx >= len(self._traces):
            return
        trace = self._traces[self._current_idx]
        spans = trace.get("spans", [])
        if not spans:
            span_widget.update("[dim]No spans[/dim]")
            return

        span_ids = [s.get("span_id") for s in spans]
        lines = ["[bold]Spans[/bold]", ""]
        in_select = self.mode == _SPAN_SELECT

        for i, span in enumerate(spans):
            parent_id = span.get("parent_id")
            is_child = parent_id and parent_id in span_ids
            kind = span.get("kind", "?")
            name = span.get("name", "unknown")
            color = KIND_COLORS.get(kind, "white")

            display_name = name if len(name) <= 24 else name[:22] + "…"

            details = []
            in_tok = span.get("input_tokens")
            out_tok = span.get("output_tokens")
            if in_tok and out_tok:
                details.append(f"{in_tok:,}→{out_tok:,} tok")
            cost = span.get("cost")
            if cost:
                details.append(f"${cost:.4f}")
            finish = span.get("finish_reason")
            if finish:
                details.append(finish)

            if in_select:
                cursor = "▸ " if i == self.span_cursor else "  "
                marker = (
                    "[red bold]◆[/red bold]"
                    if i in self._selected_spans
                    else "[dim]○[/dim]"
                )
            else:
                cursor = ""
                marker = ""

            indent = "    └─ " if is_child else ""
            kind_tag = f"[{color} bold]{kind:<5}[/{color} bold]"
            name_tag = f"[bold]{display_name}[/bold]"
            detail_str = f"  [dim]{' · '.join(details)}[/dim]" if details else ""
            line = f"{cursor}{marker} {indent}{kind_tag}  {name_tag}{detail_str}"

            if in_select and i == self.span_cursor:
                lines.append(f"[reverse]{line}[/reverse]")
            else:
                lines.append(line)

            output = span.get("output")
            if output and kind == "llm" and isinstance(output, str):
                preview = output[:80].replace("\n", " ")
                if len(output) > 80:
                    preview += "…"
                pad = "           " if is_child else "       "
                lines.append(f"{pad}[dim italic]{preview}[/dim italic]")

            input_data = span.get("input")
            if input_data and kind in ("tool", "retrieval"):
                if isinstance(input_data, dict):
                    preview = ", ".join(f"{k}={v}" for k, v in input_data.items())
                else:
                    preview = str(input_data)
                if len(preview) > 60:
                    preview = preview[:58] + "…"
                pad = "           " if is_child else "       "
                lines.append(f"{pad}[dim]→ {preview}[/dim]")

        if in_select:
            lines.append("")
            lines.append(
                "[dim]↑↓ navigate   space select   enter submit   esc cancel[/dim]"
            )

        span_widget.update("\n".join(lines))

    def watch_show_detail(self, value: bool) -> None:
        self._render_span_tree()

    def watch_mode(self, value: str) -> None:
        self._render_span_tree()

    def watch_span_cursor(self, value: int) -> None:
        self._render_span_tree()

    # ── Actions ────────────────────────────────────────────────────────────

    def action_grade_good(self) -> None:
        if self.mode != _NORMAL or not self._traces:
            return
        self._submit_grade(score=1)

    def action_grade_bad(self) -> None:
        if not self._traces:
            return
        if self.mode == _NORMAL:
            self.show_detail = True
            self.mode = _SPAN_SELECT
            self.span_cursor = 0
            self._selected_spans.clear()
            self.query_one("#span-tree", SpanSelector).focus()

    def action_skip(self) -> None:
        if self.mode != _NORMAL or not self._traces:
            return
        self._advance()

    def action_comment(self) -> None:
        if self.mode == _COMMENT or not self._traces:
            return
        self._pre_comment_mode = self.mode
        self.mode = _COMMENT
        comment_input = self.query_one("#comment-input", Input)
        comment_label = self.query_one("#comment-label", Static)
        comment_label.update("[bold]Comment[/bold]")
        comment_label.display = True
        comment_input.display = True
        comment_input.value = self._comment or ""
        comment_input.focus()

    def action_add_to_dataset(self) -> None:
        if self.mode != _NORMAL or not self._traces:
            return
        self._pre_comment_mode = self.mode
        self.mode = _ADD_DATASET
        comment_input = self.query_one("#comment-input", Input)
        comment_label = self.query_one("#comment-label", Static)
        comment_label.update("[bold]Add to Dataset (name or 'new'):[/bold]")
        comment_label.display = True
        comment_input.display = True
        comment_input.value = ""
        comment_input.focus()

    def action_toggle_detail(self) -> None:
        if self.mode == _NORMAL:
            self.show_detail = not self.show_detail

    def action_refresh(self) -> None:
        if self.mode != _NORMAL or self._picking_project:
            return
        from ..app import BanditoApp

        app: BanditoApp = self.app
        if app.project:
            self._load_grading(app.project)

    def action_toggle_filter(self) -> None:
        if self.mode != _NORMAL or self._picking_project:
            return
        from ..app import BanditoApp

        idx = _FILTER_CYCLE.index(self._filter)
        self._filter = _FILTER_CYCLE[(idx + 1) % len(_FILTER_CYCLE)]
        app: BanditoApp = self.app
        if app.project:
            self._load_grading(app.project)

    def action_go_back(self) -> None:
        if self.mode == _COMMENT:
            self._close_comment(save=False)
        elif self.mode == _ADD_DATASET:
            self._close_add_dataset()
        elif self.mode == _SPAN_SELECT:
            self.mode = _NORMAL
            self._focus_sidebar()
        else:
            self.app.pop_screen()

    def action_quit(self) -> None:
        self.app.exit()

    # ── Comment Handling ───────────────────────────────────────────────────

    @on(Input.Submitted, "#comment-input")
    def on_comment_submitted(self, event: Input.Submitted) -> None:
        if self.mode == _ADD_DATASET:
            self._handle_add_to_dataset()
        else:
            self._close_comment(save=True)

    def _handle_add_to_dataset(self) -> None:
        from ..app import BanditoApp

        if not self._traces or self._current_idx >= len(self._traces):
            self._close_add_dataset()
            return

        comment_input = self.query_one("#comment-input", Input)
        dataset_name = comment_input.value.strip()

        if not dataset_name:
            self._close_add_dataset()
            return

        trace = self._traces[self._current_idx]
        trace_id = trace["trace_id"]
        app: BanditoApp = self.app

        try:
            # Try to find existing dataset
            datasets = app.client.list_datasets(project=app.project or "")
            dataset = None
            for ds in datasets:
                if ds["name"] == dataset_name:
                    dataset = ds
                    break

            # Create new dataset if it doesn't exist
            if dataset is None:
                if dataset_name.lower() == "new":
                    # Prompt for name
                    self.query_one("#comment-label", Static).update(
                        "[bold]Enter new dataset name:[/bold]"
                    )
                    comment_input.value = ""
                    return

                new_ds = app.client.create_dataset(
                    name=dataset_name,
                    project=app.project or "",
                )
                dataset = new_ds

            # Add trace to dataset
            app.client.add_dataset_item(
                dataset_id=dataset["id"],
                trace_id=trace_id,
                source="manual",
                metadata={"added_from": "tui"},
            )

            self._notify(f"Added trace to dataset '{dataset['name']}'")

        except Exception as e:
            self._notify(f"Error: {e}")

        self._close_add_dataset()

    def _close_add_dataset(self) -> None:
        comment_input = self.query_one("#comment-input", Input)
        comment_input.display = False
        self.query_one("#comment-label", Static).display = False
        self.mode = _NORMAL
        self._focus_sidebar()

    def _notify(self, message: str) -> None:
        label = self.query_one("#trace-summary", Static)
        label.update(f"[bold green]{message}[/bold green]")

    def _close_comment(self, save: bool) -> None:
        comment_input = self.query_one("#comment-input", Input)
        if save:
            val = comment_input.value.strip()
            self._comment = val if val else None
        comment_input.display = False
        self.query_one("#comment-label", Static).display = False
        self.mode = self._pre_comment_mode
        if self.mode == _SPAN_SELECT:
            self.query_one("#span-tree", SpanSelector).focus()
        else:
            self._focus_sidebar()
        self._update_comment_display()

    def _update_comment_display(self) -> None:
        label = self.query_one("#comment-label", Static)
        if self._comment:
            label.update(
                f"[bold]Comment:[/bold] [italic]{self._comment}[/italic]"
            )
            label.display = True
        else:
            label.display = False

    # ── Key Handling ───────────────────────────────────────────────────────

    def _focus_sidebar(self) -> None:
        self.query_one("#trace-list", OptionList).focus()

    def on_key(self, event) -> None:
        if self.mode == _COMMENT:
            if event.key == "escape":
                self._close_comment(save=False)
                event.prevent_default()
            return

        if self.mode == _ADD_DATASET:
            if event.key == "escape":
                self._close_add_dataset()
                event.prevent_default()
            return

        if self.mode != _SPAN_SELECT:
            return

        if self._current_idx >= len(self._traces):
            return
        spans = self._traces[self._current_idx].get("spans", [])
        if not spans:
            return

        if event.key == "up":
            self.span_cursor = max(0, self.span_cursor - 1)
            event.prevent_default()
        elif event.key == "down":
            self.span_cursor = min(len(spans) - 1, self.span_cursor + 1)
            event.prevent_default()
        elif event.key == "space":
            if self.span_cursor in self._selected_spans:
                self._selected_spans.discard(self.span_cursor)
            else:
                self._selected_spans.add(self.span_cursor)
            self._render_span_tree()
            event.prevent_default()
        elif event.key == "enter":
            self._submit_grade(score=0)
            event.prevent_default()
        elif event.key == "escape":
            self.mode = _NORMAL
            self._focus_sidebar()
            event.prevent_default()

    # ── Grading Logic ──────────────────────────────────────────────────────

    def _submit_grade(self, score: int) -> None:
        from ..app import BanditoApp

        if not self._traces or self._current_idx >= len(self._traces):
            return
        trace = self._traces[self._current_idx]
        trace_id = trace["trace_id"]
        spans = trace.get("spans", [])
        app: BanditoApp = self.app
        comment = self._comment

        span_ids: list[str] | None = None
        if score == 0 and self._selected_spans:
            span_ids = []
            for span_idx in sorted(self._selected_spans):
                if span_idx < len(spans):
                    sid = spans[span_idx].get("span_id")
                    if sid:
                        span_ids.append(sid)

        ev_data = {
            "project": app.project or "",
            "trace_id": trace_id,
            "human_score": score,
        }
        if comment:
            ev_data["human_comment"] = comment
        if span_ids:
            ev_data["human_span_ids"] = span_ids
        app.client.write_evaluation(ev_data)

        self.mode = _NORMAL
        self._selected_spans.clear()
        self._comment = None
        self._update_comment_display()
        self._focus_sidebar()

        self._human_graded_ids.add(trace_id)

        if self._filter in (_FILTER_NEEDS_REVIEW, _FILTER_JUDGE_ONLY):
            self._remove_current_and_advance()
        else:
            self._populate_trace_list()
            self._advance()

    def _remove_current_and_advance(self) -> None:
        if not self._traces:
            return
        self._traces.pop(self._current_idx)
        self._update_title()
        self._populate_trace_list()

        if not self._traces:
            self.query_one("#trace-input", Markdown).update(
                "**All caught up!**"
            )
            self.query_one("#trace-output", Markdown).update("")
            self.query_one("#trace-summary", Static).update("")
            self.query_one("#span-tree", SpanSelector).update("")
            return

        self._current_idx = min(self._current_idx, len(self._traces) - 1)
        self.query_one("#trace-list", OptionList).highlighted = self._current_idx
        self._show_trace(self._current_idx)

    def _advance(self) -> None:
        if not self._traces:
            return
        self._current_idx = (self._current_idx + 1) % len(self._traces)
        self.query_one("#trace-list", OptionList).highlighted = self._current_idx
        self._show_trace(self._current_idx)
