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

import click
import rich
import tomli
from pydantic import ValidationError
from rich.console import Console
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(),
        console=Console(stderr=True),
    ) 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, progress)))
                    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, progress: Progress) -> None:
    child = progress.add_task(f"  pod {pod.pod_id} {pod.get_name()}", total=100)
    last_logged_decile = -1

    def on_progress(params) -> None:
        nonlocal last_logged_decile
        p = params.get("progress", 0)
        progress.update(child, completed=p)
        decile = p // 10
        if decile != last_logged_decile:
            last_logged_decile = decile
            print(f"[pod {pod.pod_id}] cloning {p}%", flush=True)

    try:
        print(f"[pod {pod.pod_id}] clone start", flush=True)
        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=pod.parent.pc_clone_specs,
            progress=on_progress,
        )
        print(f"[pod {pod.pod_id}] clone done", flush=True)
    except Exception as e:
        if sys.version_info >= (3, 11):
            e.add_note(f"pod_id={pod.pod_id}")
        raise
    finally:
        progress.remove_task(child)


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
