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

from rich import filesize
from rich.live import Live
from rich.panel import Panel
from rich.progress import Progress
from rich.progress import SpinnerColumn
from rich.progress import TextColumn
from rich.progress import BarColumn
from rich.progress import DownloadColumn
from rich.progress import TimeElapsedColumn
from rich.table import Table

from bulk_restore_tool.models import RestorableFile
from bulk_restore_tool.models import FileStatus


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

    def __init__(self, job):
        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_bar = 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_bar.add_task(
                description=description,
                total=device.totalBytes,
                completed=device.completedBytes,
            )
            device_table.add_row(device.progress_bar)
            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(Panel(device_table, height=2 + device_table.row_count))

        # extras bar (for file writers)
        self.extra: Progress = Progress(
            SpinnerColumn(), TextColumn("[progress.description]{task.description}")
        )
        self._table.add_row(self.extra)

        # 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, console=self._job._log_util._console
        )
        self._archive_tasks = {}  # type: ignore

    def start_file_progress(self, file: RestorableFile):
        task = asyncio.current_task()
        task_name = id(task)
        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}"
        task_id = self._archive_tasks.get(task_name)
        if task_id is None:
            self._archive_tasks[task_name] = device.archive_progress.add_task(
                description=description, start=False
            )
        else:
            device.archive_progress.update(task_id, description=description)

    def complete_file_progress(self, file: RestorableFile):
        device = self._job.devices[file.deviceId]
        if file.status == FileStatus.COMPLETED:
            device.progress_bar.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 complete_archive_queue(self, deviceId: str):
        task = asyncio.current_task()
        task_name = id(task)
        task_id = self._archive_tasks.get(task_name)
        if task_id is None:
            return
        device = self._job.devices[deviceId]
        device.archive_progress.remove_task(task_id)

    def get_total_progress(self):
        total_task = self._total_progress_bar.tasks[0]
        total = total_task.total
        completed = total_task.completed
        if total == 0:
            percent = 100
        else:
            percent = (completed / total) * 100
        total_size = filesize.decimal(int(total))
        completed_size = filesize.decimal(int(completed))
        return f"{completed_size} / {total_size} : {percent:.2f}%"

    def __enter__(self) -> "ProgressBar":
        if not self._job._quiet:
            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)
