from datetime import datetime
from typing import Any

from rich.console import Console
from rich.status import Status

from hil_testbench.data_structs.parameters import ParametersSchema
from hil_testbench.display.widgets.base import VisualizationMode, WidgetContext
from hil_testbench.display.widgets.parameter import ParameterValueWidget
from hil_testbench.display.widgets.task_status import TaskStatusWidget
from hil_testbench.display.widgets.trendline import TrendlineData
from hil_testbench.units.registry import PintValueFormatter


class TaskData:
    """Simple task tracking data."""

    def __init__(self, name: str) -> None:
        self.name = name
        self.status: str = "not_started"  # not_started, running, completed, failed
        self.start_time: float | None = None
        self.end_time: float | None = None
        self.error_count: int = 0
        self.parameters: dict[str, Any] = {}
        self.trendlines: dict[str, TrendlineData] = {}
        self.elapsed_time: float = 0
        self.total_duration: float = 60


class RichBackend:
    """Simple scrolling console display - shows updates as they happen."""

    def __init__(self, config: Any, logger: Any) -> None:
        self._console = Console()
        self._tasks: dict[str, TaskData] = {}
        self._status: Status | None = None
        self._schemas: dict[str, ParametersSchema] = {}
        self._mode = VisualizationMode.MINIMAL
        self._config = config
        self.logger = logger

    def initialize(self, task_names: list[str] | None = None) -> None:
        """Initialize task tracking."""
        if task_names:
            for name in task_names:
                self._tasks[name] = TaskData(name)

    def start_display(self) -> None:
        """Start the display with a spinner."""
        self._status = self._console.status(
            "[bold green]Running tasks...", spinner="dots"
        )
        self._status.start()

    def update_parameter(
        self, task_name: str, param_name: str, event: dict[str, Any], timestamp: float
    ) -> None:
        """Log parameter update immediately."""
        if task_name not in self._tasks:
            self._tasks[task_name] = TaskData(task_name)

        task = self._tasks[task_name]

        # Extract value from event
        value = event.get("value")
        if value is None:
            return

        # Store for summary
        task.parameters[param_name] = value

        # Update trendline data
        if param_name not in task.trendlines:
            task.trendlines[param_name] = TrendlineData(maxlen=10)
        task.trendlines[param_name].add(timestamp, value)

        # Update elapsed time
        if task.start_time:
            task.elapsed_time = timestamp - task.start_time

        # Try to render using widget if we have schema
        schema = self._schemas.get(task_name)
        if schema:
            field = next((f for f in schema.fields if f.name == param_name), None)
            if field:
                widget = ParameterValueWidget(value, field)
                context = WidgetContext(mode=self._mode)
                rendered = widget.render(context)
                time_str = datetime.fromtimestamp(timestamp).strftime("%H:%M:%S")
                self._console.print(
                    f"[dim]{time_str}[/dim] {task_name} | {param_name}: {rendered.rich_text}"
                )
                return

        # Fallback to old formatting
        unit = event.get("unit", "")
        precision = int(event.get("precision", 2))
        strategy = event.get("scale", event.get("scale_strategy", "raw"))
        formatter = PintValueFormatter(
            unit=unit, precision=precision, strategy=strategy
        )

        # Format and print immediately
        formatted_value = formatter.format(value)
        time_str = datetime.fromtimestamp(timestamp).strftime("%H:%M:%S")
        self._console.print(
            f"[dim]{time_str}[/dim] {task_name} | {param_name}: [cyan]{formatted_value}[/cyan]"
        )

    def update_task_status(
        self, task_name: str, status: str, start_time: float | None = None
    ) -> None:
        """Update task status."""
        if task_name not in self._tasks:
            self._tasks[task_name] = TaskData(task_name)

        task = self._tasks[task_name]
        task.status = status
        if start_time is not None:
            task.start_time = start_time
        if status in ("completed", "failed"):
            task.end_time = datetime.now().timestamp()

    def report_task_error(self, task_name: str, error: str) -> None:
        """Report an error immediately."""
        if task_name not in self._tasks:
            self._tasks[task_name] = TaskData(task_name)

        self._tasks[task_name].error_count += 1

        # Print error immediately in red
        time_str = datetime.now().strftime("%H:%M:%S")
        self._console.print(
            f"[dim]{time_str}[/dim] [bold red]ERROR[/bold red] {task_name}: {error}"
        )

    def add_log_line(self, message: str, level: str = "INFO") -> None:
        """Add a log line (shown immediately)."""
        time_str = datetime.now().strftime("%H:%M:%S")
        level_colors = {
            "DEBUG": "dim",
            "INFO": "blue",
            "WARNING": "yellow",
            "ERROR": "red",
            "CRITICAL": "bold red",
        }
        color = level_colors.get(level, "white")
        self._console.print(
            f"[dim]{time_str}[/dim] [{color}]{level}[/{color}] {message}"
        )

    def update_display(self) -> None:
        """No-op for simple backend (updates happen immediately)."""
        pass

    def shutdown(self) -> None:
        """Stop spinner and print final summary."""
        if self._status:
            self._status.stop()

        self._console.print("\n[bold]Task Summary:[/bold]")

        # Use widgets for rendering
        context = WidgetContext(mode=VisualizationMode.NORMAL)

        # Print each task with pass/fail indicator
        total_tasks = len(self._tasks)
        completed = 0
        failed = 0
        total_errors = 0

        for task in self._tasks.values():
            # Determine overall status
            if task.status == "completed" and task.error_count == 0:
                final_status = "completed"
                completed += 1
            else:
                final_status = "failed"
                failed += 1

            total_errors += task.error_count

            # Update elapsed time
            if task.start_time and task.end_time:
                task.elapsed_time = task.end_time - task.start_time

            # Get primary parameter
            primary_param = self._get_primary_param(task)

            # Get duration from config
            duration = 60  # default
            if self._config and hasattr(self._config, "duration"):
                duration = float(self._config.duration or 60)
            task.total_duration = duration

            # Build composite widget
            schema = self._schemas.get(task.name)
            trendline_data = None
            if primary_param and primary_param[0] in task.trendlines:
                trendline_data = task.trendlines[primary_param[0]]

            widget = TaskStatusWidget(
                task_name=task.name,
                status=final_status,
                primary_param=primary_param,
                trendline_data=trendline_data,
                elapsed=task.elapsed_time,
                duration=task.total_duration,
                error_count=task.error_count,
                schema=schema,
            )

            rendered = widget.render(context)
            self._console.print(rendered.rich_text)

        # Print totals
        if total_tasks > 1:
            self._console.print(
                f"\n[bold]Total:[/bold] {total_tasks} tasks | "
                f"[green]{completed} passed[/green] | "
                f"[red]{failed} failed[/red] | "
                f"[yellow]{total_errors} errors[/yellow]"
            )

    def _get_primary_param(self, task: TaskData) -> tuple[str, float] | None:
        """Get primary parameter from schema."""
        schema = self._schemas.get(task.name)
        if not schema:
            # Fallback: return first parameter
            if task.parameters:
                first_key = next(iter(task.parameters.keys()))
                return (first_key, task.parameters[first_key])
            return None

        # Find field with is_primary=True
        for field in schema.fields:
            if getattr(field, "is_primary", False):
                if field.name in task.parameters:
                    return (field.name, task.parameters[field.name])

        # Fallback: first display_enabled field
        for field in schema.fields:
            if field.display_enabled and field.name in task.parameters:
                return (field.name, task.parameters[field.name])

        # Final fallback: first parameter
        if task.parameters:
            first_key = next(iter(task.parameters.keys()))
            return (first_key, task.parameters[first_key])

        return None

    def set_schema(self, task_name: str, schema: ParametersSchema) -> None:
        """Set the schema for a task (called by pipeline)."""
        self._schemas[task_name] = schema
