from types import TracebackType
from typing import Optional, Type
from urllib.parse import urlparse

from rich.live import Live
from rich.progress import (
    Progress,
    SpinnerColumn,
    TextColumn,
    BarColumn,
    DownloadColumn,
    TimeElapsedColumn,
)
from rich.table import Table

from bulk_restore_tool.models import RestoreJob, RestorableFile


class ProgressBar:
    """Manages a group of progress bars for a restore job."""

    def __init__(self, job: RestoreJob):
        job.calculate_stats()
        self._job = job
        self._table = Table.grid()

        # progress bar for each device
        description_template = "Device guid={0} name={1}"
        justify = len(
            max(
                description_template.format(device.deviceId, device.name)
                for device in job.devices.values()
            )
        )
        for device in job.devices.values():
            device_table = Table.grid()
            device.progress = Progress(
                TextColumn("[progress.description]{task.description}"),
                BarColumn(bar_width=None),
                DownloadColumn(),
                expand=True,
            )
            description = description_template.format(
                device.deviceId, device.name
            ).ljust(justify, " ")
            device.task_id = device.progress.add_task(
                description=description,
                total=device.totalBytes,
                completed=device.completedBytes,
            )
            device_table.add_row(device.progress)
            device.archive_progress = Progress(
                TextColumn("  "),
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
            )
            device_table.add_row(device.archive_progress)
            device_table.add_row()
            self._table.add_row(device_table)

        # job total progress
        total_columns = [
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
        ]
        if len(job.devices) > 1:
            total_columns += [BarColumn(bar_width=None), DownloadColumn()]
        self._total_progress_bar = Progress(*total_columns, expand=len(job.devices) > 1)
        self._completedFiles = sum(d.completedFiles for d in job.devices.values())
        completedBytes = sum(d.completedBytes for d in job.devices.values())
        self._total_completed_task = self._total_progress_bar.add_task(
            description=f"Total files restored: [green]{self._completedFiles}[/green] Elapsed:",
            completed=completedBytes,
            total=job.totalByteCount,
        )
        self._table.add_row(self._total_progress_bar)

        # error counter
        self._error_count = 0
        self._total_errors_bar = Progress(
            TextColumn("[progress.description]{task.description}")
        )
        self._total_error_task = self._total_errors_bar.add_task(
            description=f"Total files failed: {self._error_count}"
        )
        self._table.add_row(self._total_errors_bar)

        # Live class manages refreshing progress bars together
        self._live = Live(self._table, refresh_per_second=10)
        self._file_tasks = {}  # type: ignore

    def start_file_progress(self, file: RestorableFile):
        device = self._job.devices[file.deviceId]
        node = urlparse(file.storageNodeWebUrl).netloc.split(":")[0]
        details = f"fileId={file.fileId} from archive={file.archiveId} @ {file.destinationName} - {node}, size={file.size}"
        if file.failure_count > 0:
            description = f"Retrying {details} - [red]retry {file.failure_count}[/red]"
        else:
            description = f"Restoring {details}"
        self._file_tasks[file.fileId] = device.archive_progress.add_task(
            description=description, start=False
        )

    def complete_file_progress(self, file: RestorableFile):
        device = self._job.devices[file.deviceId]
        device.archive_progress.remove_task(self._file_tasks[file.fileId])
        if file.completed:
            device.progress.update(device.task_id, advance=file.sourceLengthBytes)
            self._completedFiles += 1
            self._total_progress_bar.update(
                task_id=self._total_completed_task,
                description=f"Total files restored: [green]{self._completedFiles}[/green] Elapsed:",
                advance=file.sourceLengthBytes,
            )
        else:
            if file.failure_count > 2:
                self._error_count += 1
                self._total_errors_bar.update(
                    task_id=self._total_error_task,
                    description=f"Total files failed: [red]{self._error_count}",
                )

    def __enter__(self) -> "ProgressBar":
        self._live.start()
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        self._live.stop()

    def print(self, *args, **kwargs):
        self._live.console.print(*args, **kwargs)
