import logging
import sys
from pathlib import Path
from textwrap import indent

import click
import rich
import tomli
from pydantic import ValidationError
from rich.markup import escape as e
from rich.progress import MofNCompleteColumn, Progress

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

from netlab.enums import VirtualMachineRole

from ..async_client import NetlabConnection
from ..enums import PodState, RemoveVMS
from . import schema
from .compiler import (
    CompiledDelete,
    CompiledManifest,
    CompiledPod,
    CompiledRange,
    CompiledRangeErrors,
)
from .fetcher import _fetch_hosts, _fetch_pods
from .types import ExistingPod

log = logging.getLogger(__name__)


async def main(manifest_paths: list[Path], auto_confirm: bool) -> None:
    found_errors = False

    manifests = []

    for manifest_path in manifest_paths:
        log.debug("Loading manifest", extra={"path": manifest_path})

        try:
            manifest = schema.Manifest.model_validate(tomli.loads(manifest_path.read_text()))
            manifests.append(manifest)
        except ValidationError as validation_error:
            print(validation_error)
            rich.print("[red bold]Found errors. Cannot continue.")
            return
            # found_errors |= True

    from .compiler import client_factory_from_manifest

    async with client_factory_from_manifest(manifests[0]) as client:
        pods = await _fetch_pods(client, manifests[0])
        hosts = await _fetch_hosts(client)

    compiled_manifest = CompiledManifest.from_schema(manifests[0], manifest_paths[0], hosts, pods)

    for i, r in enumerate(compiled_manifest.ranges):
        print(f"{compiled_manifest.manifest_path}#{i}")

        if isinstance(r, CompiledRange | CompiledDelete):
            for pod_clo in r.pods:
                conflict_kind = pod_clo.pod_op()
                print(indent(str(conflict_kind), "  "))
        else:
            found_errors |= True
            for error_code, error_str in r.errors:
                rich.print(f"  [red]{e(error_code):20} {e(error_str)}")

    if found_errors:
        rich.print("[red bold]Found errors. Cannot continue.")
        return

    if auto_confirm is False:
        click.confirm("Continue?", abort=True)

    await invoke_manifest(compiled_manifest)


async def invoke_manifest(compiled_manifest: CompiledManifest) -> None:
    pods_to_delete: list[ExistingPod] = []
    pods_to_create: list[CompiledPod] = []
    pods_to_updown: list[CompiledPod] = []

    for r in compiled_manifest.ranges:
        assert not isinstance(r, CompiledRangeErrors)
        for pod in r.pods:
            op = pod.pod_op()

            create, delete = op.op()

            if create:
                pods_to_create.append(create)
            if delete:
                pods_to_delete.append(delete)
                #         pods_to_delete.append(existing_pod)
            # match op:
            #     case pod_op.Absent() | pod_op.Superceded() | pod_op.Present():
            #         pass
            #     case pod_op.Rebuild(pod=pod, existing=existing_pod):
            #         pods_to_create.append(pod)
            #         pods_to_delete.append(existing_pod)
            #     case pod_op.Create(pod=pod):
            #         pods_to_create.append(pod)
            #     case pod_op.Delete(existing=existing_pod)
            #         pods_to_delete.append(existing_pod)

            if isinstance(pod.parent, CompiledRange) and pod.parent.get_role() is not VirtualMachineRole.TEMPLATE:
                pods_to_updown.append(pod)

            # if pod.pod_op() in ("REBUILD", "CREATE"):
            #     pods_to_create.append(pod)
            # if (
            #     pod.pod_op() in ("PRESENT", "REBUILD", "CREATE")
            #     and pod.get_role() is not VirtualMachineRole.TEMPLATE
            # ):
            #     pods_to_updown.append(pod)

    with Progress(
        MofNCompleteColumn(),
        *Progress.get_default_columns(),
    ) as progress:
        async with compiled_manifest.client as client:
            delete_progress = progress.add_task("Deleting pods", total=len(pods_to_delete))
            create_progress = progress.add_task("Creating pods", total=len(pods_to_create))
            updown_progress = progress.add_task("Setting pods", total=len(pods_to_updown))

            throt = compiled_manifest.throttle_delete.make()
            async with TaskGroup() as tg:
                for delpod in reversed(pods_to_delete):
                    bucket = await throt.bucket()
                    task = tg.create_task(bucket(pod_deleter(client, delpod)))
                    task.add_done_callback(lambda f: progress.advance(delete_progress, 1))

            throt = compiled_manifest.throttle_create.make()
            async with TaskGroup() as tg:
                for pod in pods_to_create:
                    bucket = await throt.bucket()
                    task = tg.create_task(bucket(pod_creator(client, pod)))
                    task.add_done_callback(lambda f: progress.advance(create_progress, 1))

            throt = compiled_manifest.throttle_updown.make()
            async with TaskGroup() as tg:
                for pod in pods_to_updown:
                    bucket = await throt.bucket()
                    task = tg.create_task(bucket(pod_onliner(client, pod)))
                    task.add_done_callback(lambda f: progress.advance(updown_progress, 1))


async def pod_creator(client: NetlabConnection, pod: CompiledPod) -> None:
    try:
        print(f"cloning pod {pod.pod_id}")
        pc_clone_specs = pod.parent.pc_clone_specs
        await client.pod_clone_task(
            source_pod_id=pod.parent.source_pod.pod_id,
            clone_pod_id=pod.pod_id,
            clone_pod_name=pod.get_name(),
            pc_clone_specs=pc_clone_specs,
        )

        print(f"done cloning pod {pod.pod_id}")
    except Exception as e:
        if sys.version_info >= (3, 11):
            e.add_note(f"pod_id={pod.pod_id}")
        raise


async def pod_deleter(client: NetlabConnection, pod: ExistingPod) -> None:
    try:
        print(f"Removing pod {pod.pod_id}")
        await client.pod_remove_task(
            pod_id=pod.pod_id,
            remove_vms=RemoveVMS.DISK,
        )

        print(f"Done removing pod {pod.pod_id}")
    except Exception as e:
        if sys.version_info >= (3, 11):
            e.add_note(f"pod_id={pod.pod_id}")
        raise


async def pod_onliner(client: NetlabConnection, pod: CompiledPod) -> None:
    try:
        state = PodState.ONLINE if pod.parent.pod_online else PodState.OFFLINE
        print(f"Setting pod {pod.pod_id} state={state}")
        assert pod.parent
        await client.pod_state_change(
            pod_id=pod.pod_id,
            state=state,
        )
    except Exception as e:
        if sys.version_info >= (3, 11):
            e.add_note(f"pod_id={pod.pod_id}")
        raise
