from __future__ import annotations

import random
from pathlib import Path
from typing import Any, Dict, Tuple

import yaml
from rich.text import Text
from textual.app import ComposeResult
from textual.containers import Container, Horizontal
from textual.screen import Screen
from textual.widgets import (
    Button,
    Footer,
    Header,
    Input,
    LoadingIndicator,
    Static,
    Tree,
)
from textual.binding import Binding

from crystallize.experiments.experiment import Experiment
from crystallize.experiments.experiment_graph import ExperimentGraph

from ..constants import ASCII_ART_ARRAY
from ..discovery import discover_configs
from ..screens.create_experiment import CreateExperimentScreen
from ..screens.run import _launch_run
from ..utils import update_replicates


class ExperimentTree(Tree):
    """Tree widget with custom binding for experiment selection."""

    BINDINGS = [b for b in Tree.BINDINGS if getattr(b, "key", "") != "enter"] + [
        Binding("enter", "run_selected", "Run", show=True)
    ]

    def action_run_selected(self) -> None:  # pragma: no cover - delegates
        screen = self.screen
        if screen is not None and hasattr(screen, "action_run_selected"):
            screen.action_run_selected()

        try:
            line = self._tree_lines[self.cursor_line]
        except IndexError:
            pass
        else:
            node = line.path[-1]
            self.post_message(Tree.NodeSelected(node))


class SelectionScreen(Screen):
    """Main screen for selecting experiments or graphs."""

    BINDINGS = [
        ("q", "quit", "Quit"),
        ("r", "refresh", "Refresh"),
        ("c", "create_experiment", "Create Experiment"),
        ("e", "show_errors", "Errors"),
    ]

    def __init__(self) -> None:
        super().__init__()
        self._load_errors: Dict[str, BaseException] = {}
        self._experiments: Dict[str, Dict[str, Any]] = {}
        self._graphs: Dict[str, Dict[str, Any]] = {}
        self._selected_obj: Dict[str, Any] | None = None

    def _update_details(self, data: Dict[str, Any]) -> None:
        """Populate the details panel with information from ``data``."""

        details = self.query_one("#details", Static)
        info = yaml.safe_load(Path(data["path"]).read_text()) or {}

        desc = info.get("description", data.get("doc", ""))
        repl = info.get("replicates", data.get("replicates", 1))
        steps = info.get("steps", [])
        treatments = list((info.get("treatments") or {}).keys())
        outputs = [
            v.get("file_name", k) if isinstance(v, dict) else str(v)
            for k, v in (info.get("outputs") or {}).items()
        ]

        details_text = [
            f"[bold]Type: {data['type']}[/bold]",
            "",
            desc,
            f"Replicates: {repl}",
        ]

        if steps:
            details_text.append("[bold]Steps:[/bold]")
            details_text.extend(f"- {s}" for s in steps)

        if treatments:
            details_text.append("[bold]Treatments:[/bold]")
            details_text.extend(f"- {t}" for t in treatments)

        if outputs:
            details_text.append("[bold]Outputs:[/bold]")
            details_text.extend(f"- {o}" for o in outputs)

        details.update("\n".join(details_text))

        rep_input = self.query_one("#replicate-input", Input)
        rep_input.value = str(repl)

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        yield Static(random.choice(ASCII_ART_ARRAY), id="title")
        with Container(id="main-container"):
            yield LoadingIndicator()
            yield Static(
                "Scanning for experiments and graphs...",
                id="loading-text",
            )
        yield Footer()

    async def on_mount(self) -> None:
        self.run_worker(self._discover())

    def action_refresh(self) -> None:
        self.run_worker(self._discover())

    def action_create_experiment(self) -> None:
        self.app.push_screen(CreateExperimentScreen())

    def _discover_sync(
        self,
    ) -> Tuple[
        Dict[str, Dict[str, Any]],
        Dict[str, Dict[str, Any]],
        Dict[str, BaseException],
    ]:
        """Locate ``config.yaml`` files and classify them."""

        return discover_configs(Path("."))

    async def _discover(self) -> None:
        worker = self.run_worker(self._discover_sync, thread=True)
        graphs, experiments, errors = await worker.wait()
        self._load_errors = errors
        self._experiments = experiments
        self._graphs = graphs

        main_container = self.query_one("#main-container")
        await main_container.remove_children()

        await main_container.mount(Static("Select an object to run:"))

        horizontal = Horizontal()
        await main_container.mount(horizontal)

        left_panel = Container(classes="left-panel")
        await horizontal.mount(left_panel)

        tree = ExperimentTree("root", id="object-tree")
        tree.show_root = False
        await left_panel.mount(tree)

        groups: dict[str, list[tuple[str, Dict[str, Any]]]] = {}
        for label, info in graphs.items():
            groups.setdefault(info["cli"]["group"], []).append(("Graph", info))
        for label, info in experiments.items():
            groups.setdefault(info["cli"]["group"], []).append(("Experiment", info))

        for group_name in sorted(groups):
            parent = tree.root.add(group_name, expand=True)
            items = sorted(groups[group_name], key=lambda t: t[1]["cli"]["priority"])
            for obj_type, info in items:
                label = info["label"]
                icon = info["cli"]["icon"]
                color = info["cli"].get("color")
                text = Text(f"{icon} {label} ({obj_type})")
                if color:
                    text.stylize(color)
                parent.add_leaf(
                    text,
                    {
                        "path": info["path"],
                        "label": label,
                        "type": obj_type,
                        "doc": info["description"] or "No description available.",
                    },
                )

        right_panel = Container(classes="right-panel")
        await horizontal.mount(right_panel)
        await right_panel.mount(Static(id="details", classes="details-panel"))

        btn_container = Container(id="select-button-container")
        await right_panel.mount(btn_container)
        await btn_container.mount(Button("Run", id="run-btn"))
        await btn_container.mount(Input(id="replicate-input", placeholder="replicates"))

        if self._load_errors:
            await main_container.mount(
                Static(
                    f"Failed to load {len(self._load_errors)} file(s), press e for more details",
                    id="error-msg",
                )
            )

        tree.focus()

    async def _run_interactive_and_exit(self, info: Dict[str, Any]) -> None:
        cfg = info["path"]
        obj_type = info["type"]
        try:
            if obj_type == "Graph":
                obj = ExperimentGraph.from_yaml(cfg)
            else:
                obj = Experiment.from_yaml(cfg)
        except BaseException as exc:  # noqa: BLE001
            self._load_errors[str(cfg)] = exc
            from ..screens.load_errors import LoadErrorsScreen

            self.app.push_screen(LoadErrorsScreen({str(cfg): exc}))
            return

        await _launch_run(self.app, obj)

    def action_run_selected(self) -> None:
        if self._selected_obj is not None:
            self.run_worker(self._run_interactive_and_exit(self._selected_obj))

    async def on_tree_node_highlighted(self, event: Tree.NodeHighlighted) -> None:
        if event.node.data is not None:
            data = event.node.data
            self._update_details(data)
            self._selected_obj = data
        else:
            details = self.query_one("#details", Static)
            details.update("")
            self._selected_obj = None

    async def on_tree_node_selected(self, event: Tree.NodeSelected) -> None:
        if event.node.data is not None:
            data = event.node.data
            self._update_details(data)
            self._selected_obj = data

    def action_show_errors(self) -> None:
        if self._load_errors:
            from ..screens.load_errors import LoadErrorsScreen

            self.app.push_screen(LoadErrorsScreen(self._load_errors))

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "run-btn":
            self.action_run_selected()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        if event.input.id == "replicate-input" and self._selected_obj is not None:
            try:
                new_val = int(event.value)
            except ValueError:
                self.app.bell()
                return
            update_replicates(Path(self._selected_obj["path"]), new_val)
            self._selected_obj["replicates"] = new_val
            self._update_details(self._selected_obj)
