# -*- coding: utf-8 -*-
"""Location: ./tests/load/utils/progress.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Progress tracking utilities with rich display.
"""

import sys
import time
from contextlib import contextmanager
from typing import Dict, Optional

from rich.console import Console, RenderableType
from rich.live import Live
from rich.panel import Panel
from rich.progress import (
    Progress,
    SpinnerColumn,
    BarColumn,
    TextColumn,
    TimeRemainingColumn,
    TimeElapsedColumn,
    MofNCompleteColumn,
    TaskProgressColumn,
)
from rich.table import Table
from rich.layout import Layout
from rich.text import Text
from collections import deque


class ProgressTracker:
    """Track progress of data generation with rich display and statistics."""

    def __init__(self, total: int, desc: str, unit: str = "records"):
        """Initialize progress tracker.

        Args:
            total: Total number of items to process
            desc: Description of the operation
            unit: Unit name for progress display
        """
        self.total = total
        self.desc = desc
        self.unit = unit
        self.start_time = time.time()
        self.current = 0
        self.task_id = None
        self.progress = None

    def __enter__(self):
        """Context manager entry."""
        self.progress = Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description}"),
            BarColumn(complete_style="green", finished_style="bold green"),
            TaskProgressColumn(),
            MofNCompleteColumn(),
            TextColumn("•"),
            TimeElapsedColumn(),
            TextColumn("•"),
            TimeRemainingColumn(),
            TextColumn("[cyan]{task.fields[rate]}/s"),
        )
        self.task_id = self.progress.add_task(
            self.desc,
            total=self.total,
            rate="0"
        )
        self.progress.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        if self.progress:
            self.progress.stop()

    def update(self, n: int = 1, **kwargs):
        """Update progress by n items.

        Args:
            n: Number of items to add to progress
            **kwargs: Additional fields to update
        """
        if self.progress and self.task_id is not None:
            self.current += n
            rate = self.get_rate()
            self.progress.update(
                self.task_id,
                advance=n,
                rate=f"{rate:,.0f}",
                **kwargs
            )

    def get_elapsed_time(self) -> float:
        """Get elapsed time in seconds.

        Returns:
            Elapsed time
        """
        return time.time() - self.start_time

    def get_rate(self) -> float:
        """Get processing rate (items per second).

        Returns:
            Items per second
        """
        elapsed = self.get_elapsed_time()
        if elapsed > 0:
            return self.current / elapsed
        return 0.0


class MultiProgressTracker:
    """Track multiple progress bars simultaneously with rich live display."""

    def __init__(self, console: Optional[Console] = None, max_log_lines: int = 3):
        """Initialize multi-progress tracker.

        Args:
            console: Rich console instance (creates new if None)
            max_log_lines: Maximum number of log lines to display (default: 3)
        """
        self.console = console or Console()
        self.progress = Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description:30}"),
            BarColumn(complete_style="green", finished_style="bold green"),
            TaskProgressColumn(),
            MofNCompleteColumn(),
            TextColumn("•"),
            TimeElapsedColumn(),
            TextColumn("•"),
            TimeRemainingColumn(),
            TextColumn("[cyan]{task.fields[rate]}/s"),
            expand=False,
        )
        self.tasks: Dict[str, int] = {}
        self.stats: Dict[str, Dict] = {}
        self.live: Optional[Live] = None
        self.layout = Layout()
        self.start_time = time.time()
        self.total_records = 0
        self.total_completed = 0
        self.max_log_lines = max_log_lines
        self.log_buffer: deque = deque(maxlen=max_log_lines)
        self.show_logs = True
        self.is_interactive = sys.stdout.isatty()
        self.last_status_print = 0

    def add_task(self, name: str, total: int, desc: str):
        """Add a new progress task.

        Args:
            name: Unique name for the task
            total: Total items
            desc: Description
        """
        task_id = self.progress.add_task(
            desc,
            total=total,
            rate="0",
            visible=False
        )
        self.tasks[name] = task_id
        self.stats[name] = {
            "total": total,
            "completed": 0,
            "start_time": None,
            "end_time": None,
            "rate": 0.0,
        }
        self.total_records += total

    def start_task(self, name: str):
        """Start a task and make it visible.

        Args:
            name: Task name
        """
        if name in self.tasks:
            self.progress.update(self.tasks[name], visible=True)
            self.stats[name]["start_time"] = time.time()

            # For non-interactive terminals, print start message
            if not self.is_interactive:
                total = self.stats[name]["total"]
                self.console.print(f"[yellow]⚡[/yellow] Starting [cyan]{name}[/cyan]: {total:,} records")

    def update(self, name: str, n: int = 1):
        """Update a specific task.

        Args:
            name: Task name
            n: Number of items
        """
        if name in self.tasks:
            task_id = self.tasks[name]
            self.stats[name]["completed"] += n
            self.total_completed += n

            # Calculate rate
            elapsed = time.time() - (self.stats[name]["start_time"] or time.time())
            if elapsed > 0:
                rate = self.stats[name]["completed"] / elapsed
                self.stats[name]["rate"] = rate
            else:
                rate = 0.0

            self.progress.update(
                task_id,
                advance=n,
                rate=f"{rate:,.0f}"
            )

    def complete_task(self, name: str):
        """Mark a task as complete.

        Args:
            name: Task name
        """
        if name in self.tasks:
            # Ensure completed count matches total
            total = self.stats[name]["total"]
            current_completed = self.stats[name].get("completed", 0)

            # Update total_completed with any remaining records
            remaining = total - current_completed
            if remaining > 0:
                self.total_completed += remaining

            self.stats[name]["completed"] = total  # Always set to total when completing
            self.stats[name]["end_time"] = time.time()
            self.progress.update(self.tasks[name], completed=total)

            # For non-interactive terminals, print completion message
            if not self.is_interactive:
                rate = self.stats[name].get("rate", 0)
                self.console.print(f"[green]✓[/green] Completed [cyan]{name}[/cyan]: {total:,} records ([cyan]{rate:,.0f}/s[/cyan])")

    def log(self, message: str, style: str = "dim"):
        """Add a log message to the scrolling log panel.

        Args:
            message: Log message to display
            style: Rich style for the message
        """
        timestamp = time.strftime("%H:%M:%S")
        styled_message = f"[dim]{timestamp}[/dim] {message}"
        self.log_buffer.append((styled_message, style))

    def _make_log_panel(self) -> RenderableType:
        """Create scrolling log panel.

        Returns:
            Rich renderable with recent log messages
        """
        if not self.log_buffer:
            return Text("No log messages yet...", style="dim italic")

        # Create a group of text objects to render markup properly
        from rich.console import Group
        log_lines = []
        for message, style in self.log_buffer:
            # Use console.render_str to properly parse markup
            log_lines.append(Text.from_markup(message))

        return Group(*log_lines)

    def _make_stats_table(self) -> Table:
        """Create overall statistics table.

        Returns:
            Rich table with overall statistics
        """
        table = Table(show_header=False, box=None, padding=(0, 1), collapse_padding=True)
        table.add_column("Metric", style="bold cyan", no_wrap=True, width=18)
        table.add_column("Value", style="bold white")

        elapsed = time.time() - self.start_time
        overall_rate = self.total_completed / elapsed if elapsed > 0 else 0

        table.add_row("Total Records", f"{self.total_records:,}")
        table.add_row("Completed", f"[green]{self.total_completed:,}[/green]")
        table.add_row("Remaining", f"[yellow]{self.total_records - self.total_completed:,}[/yellow]")
        table.add_row("Overall Rate", f"[cyan]{overall_rate:,.0f} records/s[/cyan]")
        table.add_row("Elapsed Time", f"{elapsed:.1f}s")

        if overall_rate > 0 and self.total_completed < self.total_records:
            eta = (self.total_records - self.total_completed) / overall_rate
            table.add_row("ETA", f"[magenta]{eta:.1f}s[/magenta]")

        return table

    def _make_generator_status_table(self) -> Table:
        """Create detailed generator status table.

        Returns:
            Rich table with per-generator status
        """
        table = Table(show_header=True, box=None, padding=(0, 1), expand=True, show_lines=False)
        table.add_column("Generator", style="bold", no_wrap=True, width=30)
        table.add_column("Status", style="bold", width=13)
        table.add_column("Progress", justify="right", width=20)
        table.add_column("Rate", justify="right", width=13)

        # Categorize generators
        completed = []
        in_progress = []
        pending = []

        for name in self.tasks.keys():
            if name not in self.stats:
                continue  # Skip if stats not initialized

            task_info = self.stats[name]
            total = task_info.get("total", 0)
            current = task_info.get("completed", 0)

            if current >= total and total > 0:
                status = "completed"
                completed.append((name, task_info))
            elif current > 0:
                status = "in_progress"
                in_progress.append((name, task_info))
            else:
                status = "pending"
                pending.append((name, task_info))

        # Add summary row with counts
        total_generators = len(self.tasks)
        table.add_row(
            f"[bold]Summary: {total_generators} generators[/bold]",
            f"[green]✓ {len(completed)}[/green] [yellow]⚡ {len(in_progress)}[/yellow] [dim]⏳ {len(pending)}[/dim]",
            "",
            "",
        )

        # Only add section separator if there are generators to show
        if completed or in_progress or pending:
            table.add_section()

        # Show in-progress generators first
        for name, info in in_progress:
            rate = info.get("rate", 0)
            current = info.get("completed", 0)
            total = info.get("total", 0)
            pct = (current / total * 100) if total > 0 else 0

            table.add_row(
                name,
                "[yellow]⚡ Active[/yellow]",
                f"[yellow]{current:,}/{total:,} ({pct:.0f}%)[/yellow]",
                f"[cyan]{rate:,.0f}/s[/cyan]"
            )

        # Show completed generators
        for name, info in completed:
            rate = info.get("rate", 0)
            total = info.get("total", 0)
            current = info.get("completed", 0)

            table.add_row(
                name,
                "[green]✓ Done[/green]",
                f"[green]{current:,}/{total:,} (100%)[/green]",
                f"[dim]{rate:,.0f}/s[/dim]"
            )

        # Show pending generators
        for name, info in pending:
            total = info.get("total", 0)

            table.add_row(
                name,
                "[dim]⏳ Pending[/dim]",
                f"[dim]0/{total:,} (0%)[/dim]",
                "[dim]—[/dim]"
            )

        return table

    def _make_layout(self) -> Layout:
        """Create layout with progress and stats.

        Returns:
            Rich layout
        """
        layout = Layout()

        if self.show_logs:
            # Compact layout: stats(7) + generators(flex) + progress(8) + logs(5)
            layout.split_column(
                Layout(name="stats", size=7),
                Layout(name="generators"),  # Takes remaining space
                Layout(name="progress", size=8),
                Layout(name="logs", size=self.max_log_lines + 2),  # Compact: 3 logs + 2 for border
            )
        else:
            layout.split_column(
                Layout(name="stats", size=7),
                Layout(name="generators"),  # Takes remaining space
                Layout(name="progress", size=8),
            )

        # Overall stats panel
        stats_table = self._make_stats_table()
        layout["stats"].update(
            Panel(
                stats_table,
                title="[bold]Overall Statistics[/bold]",
                border_style="cyan",
            )
        )

        # Generator status table
        generator_table = self._make_generator_status_table()
        layout["generators"].update(
            Panel(
                generator_table,
                title="[bold]Generator Status (Detailed)[/bold]",
                border_style="blue",
            )
        )

        # Progress bars (only show active ones)
        layout["progress"].update(self.progress)

        # Log panel (if enabled)
        if self.show_logs:
            log_panel = self._make_log_panel()
            layout["logs"].update(
                Panel(
                    log_panel,
                    title="[bold]Activity Log[/bold]",
                    border_style="green",
                )
            )

        return layout

    @contextmanager
    def live_display(self):
        """Context manager for live display.

        Yields:
            Self for chaining
        """
        try:
            if self.is_interactive:
                # Use Rich Live display for interactive terminals
                with Live(
                    self._make_layout(),
                    console=self.console,
                    refresh_per_second=10,  # Increased from 4 to 10 for more responsive updates
                    transient=False,
                    auto_refresh=True,  # Force auto-refresh
                ) as live:
                    self.live = live
                    yield self
            else:
                # For non-interactive terminals (piped output), use simple console prints
                self.console.print("[bold cyan]Starting data generation...[/bold cyan]")
                self.console.print(f"[dim]Total generators: {len(self.tasks)}[/dim]")
                self.console.print(f"[dim]Total records: {self.total_records:,}[/dim]")
                self.console.print()
                yield self
                # Print final summary
                self.console.print()
                self._print_final_summary()
        finally:
            self.live = None

    def refresh(self):
        """Refresh the live display."""
        if self.live:
            self.live.update(self._make_layout())
        elif not self.is_interactive:
            # For non-interactive terminals, print periodic status updates
            current_time = time.time()
            if current_time - self.last_status_print >= 5.0:  # Print every 5 seconds
                self._print_status_update()
                self.last_status_print = current_time

    def _print_status_update(self):
        """Print a status update to console (for non-interactive terminals)."""
        # Count generators by status
        completed_count = 0
        in_progress_count = 0
        pending_count = 0

        for name in self.tasks.keys():
            if name not in self.stats:
                continue
            task_info = self.stats[name]
            total = task_info.get("total", 0)
            current = task_info.get("completed", 0)

            if current >= total and total > 0:
                completed_count += 1
            elif current > 0:
                in_progress_count += 1
            else:
                pending_count += 1

        elapsed = time.time() - self.start_time
        overall_rate = self.total_completed / elapsed if elapsed > 0 else 0
        pct = (self.total_completed / self.total_records * 100) if self.total_records > 0 else 0

        self.console.print(
            f"[cyan]Progress:[/cyan] {self.total_completed:,}/{self.total_records:,} ({pct:.1f}%) | "
            f"[green]✓ {completed_count}[/green] [yellow]⚡ {in_progress_count}[/yellow] [dim]⏳ {pending_count}[/dim] | "
            f"[cyan]{overall_rate:,.0f} rec/s[/cyan]"
        )

    def _print_final_summary(self):
        """Print final summary for non-interactive terminals."""
        elapsed = time.time() - self.start_time
        overall_rate = self.total_completed / elapsed if elapsed > 0 else 0

        self.console.print("[bold green]Generation Complete![/bold green]")
        self.console.print(f"Total Records: {self.total_completed:,}/{self.total_records:,}")
        self.console.print(f"Duration: {elapsed:.2f}s")
        self.console.print(f"Overall Rate: {overall_rate:,.0f} records/s")
        self.console.print()

        # Print generator breakdown
        table = Table(show_header=True, box=None, padding=(0, 1))
        table.add_column("Generator", style="cyan", no_wrap=True)
        table.add_column("Records", justify="right", style="green")
        table.add_column("Rate", justify="right", style="yellow")

        for name in self.tasks.keys():
            if name not in self.stats:
                continue
            task_info = self.stats[name]
            total = task_info.get("total", 0)
            rate = task_info.get("rate", 0)
            table.add_row(name, f"{total:,}", f"{rate:,.0f}/s")

        self.console.print(table)

    def close_all(self):
        """Close all progress trackers."""
        for name in self.tasks:
            self.complete_task(name)
        if self.live:
            self.live.stop()


class SimpleProgressTracker:
    """Simple progress tracker for single operations."""

    def __init__(self, console: Optional[Console] = None):
        """Initialize simple progress tracker.

        Args:
            console: Rich console instance
        """
        self.console = console or Console()
        self.start_time = time.time()

    def print_step(self, message: str, style: str = "bold blue"):
        """Print a step message.

        Args:
            message: Message to print
            style: Rich style string
        """
        self.console.print(f"[{style}]→[/{style}] {message}")

    def print_success(self, message: str):
        """Print success message.

        Args:
            message: Message to print
        """
        self.console.print(f"[bold green]✓[/bold green] {message}")

    def print_error(self, message: str):
        """Print error message.

        Args:
            message: Message to print
        """
        self.console.print(f"[bold red]✗[/bold red] {message}")

    def print_warning(self, message: str):
        """Print warning message.

        Args:
            message: Message to print
        """
        self.console.print(f"[bold yellow]⚠[/bold yellow] {message}")

    def print_info(self, message: str):
        """Print info message.

        Args:
            message: Message to print
        """
        self.console.print(f"[cyan]ℹ[/cyan] {message}")

    def print_stats(self, stats: Dict):
        """Print statistics in a formatted table.

        Args:
            stats: Dictionary of statistics
        """
        table = Table(show_header=True, box=None)
        table.add_column("Generator", style="cyan", no_wrap=True)
        table.add_column("Generated", justify="right", style="green")
        table.add_column("Inserted", justify="right", style="yellow")
        table.add_column("Time", justify="right", style="blue")
        table.add_column("Rate", justify="right", style="magenta")

        for name, data in stats.items():
            if isinstance(data, dict):
                generated = data.get("generated", 0)
                inserted = data.get("inserted", 0)
                duration = data.get("duration", 0)
                rate = generated / duration if duration > 0 else 0

                table.add_row(
                    name,
                    f"{generated:,}",
                    f"{inserted:,}",
                    f"{duration:.2f}s",
                    f"{rate:,.0f}/s"
                )

        self.console.print(table)

    def print_summary(self, total_records: int, duration: float, profile: str):
        """Print final summary.

        Args:
            total_records: Total records generated
            duration: Total duration in seconds
            profile: Profile name
        """
        rate = total_records / duration if duration > 0 else 0

        panel = Panel(
            f"[bold]Profile:[/bold] {profile}\n"
            f"[bold]Total Records:[/bold] {total_records:,}\n"
            f"[bold]Duration:[/bold] {duration:.2f}s\n"
            f"[bold]Rate:[/bold] {rate:,.2f} records/second",
            title="[bold green]Generation Complete[/bold green]",
            border_style="green",
        )
        self.console.print(panel)