import sys

from rich.progress import Progress, TaskID

if sys.version_info >= (3, 11):
    from asyncio import TaskGroup
else:
    from taskgroup import TaskGroup

from ..async_client import NetlabConnection
from . import schema
from .types import ExistingPod, NetlabHost


async def _fetch_hosts(client: NetlabConnection) -> dict[int | str, NetlabHost]:
    host_map: dict[int | str, NetlabHost] = {}

    for host_raw in await client.vm_host_list():
        host = NetlabHost(
            host_raw["vh_id"],
            host_raw["vh_name"],
        )
        host_map[host.vh_id] = host
        host_map[host.vh_name] = host

    return host_map


async def _fetch_pods(client: NetlabConnection, manifest: schema.Manifest) -> dict[int, ExistingPod]:
    result: dict[int, ExistingPod] = {}
    existing_pod_ids = set(await client.pod_list_used_ids())
    created_pod_ids: set[int] = set()
    source_pod_ids = set()
    for r in manifest.range:
        if r.delete:
            continue
        rng = set(r.get_range())
        source_pod_ids.add(r.source_pod_id)
        assert rng.isdisjoint(created_pod_ids)
        created_pod_ids.update(rng)

    async def inner(pod_id: int, progress_task: TaskID, progress: Progress) -> None:
        try:
            raw_pod = await client.pod_get(pod_id=pod_id, properties="all")
            # progress.console.print(f"Fetching pod {pod_id}")

            try:
                host = NetlabHost(
                    raw_pod["remote_pc"][0]["vh_id"],
                    raw_pod["remote_pc"][0]["vh_name"],
                )
                vm_role = raw_pod["remote_pc"][0]["vm_role"]
            except KeyError:
                host = None
                vm_role = None

            existing_pod = ExistingPod(
                pod_id=pod_id,
                current_host=host,
                current_role=vm_role,
                current_name=raw_pod["pod_name"],
                current_pod_type=raw_pod["pt_id"],
                pt_org=raw_pod["pt_org"],
                pt_apdid=raw_pod["pt_apdid"],
            )
            result[pod_id] = existing_pod
            progress.update(progress_task, advance=1)
        except Exception as e:
            if sys.version_info >= (3, 11):
                e.add_note(f"Fetching pod_id={pod_id}")
            raise e

    throt = manifest.throttles.fetch.make()
    pod_ids = existing_pod_ids.intersection(created_pod_ids | source_pod_ids)
    with Progress() as progress:
        progress_task = progress.add_task("Fetching pods", total=len(pod_ids))

        async with TaskGroup() as tg:
            for pod_id in pod_ids:
                bucket = await throt.bucket()
                tg.create_task(bucket(inner(pod_id, progress_task, progress)))

    return result
