import asyncio
import contextlib
import io
import tarfile
import time
from collections.abc import Iterable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from functools import cache
from http import HTTPStatus
from pathlib import Path
from typing import Any

import docker
import yaml
from kubernetes.client import ApiClient, ApiException
from kubernetes.dynamic import DynamicClient
from kubernetes.dynamic.exceptions import ResourceNotFoundError
from loguru import logger

from kona.schema.models import KonaChallengeConfig, KonaGlobalConfig
from kona.util.jinja import render_template, render_template_values

from .kubernetes import load_kubeconfig, resolve_cluster_names


DEFAULT_NAMESPACE = 'default'


@dataclass(frozen=True)
class BuiltDockerImage:
    path: Path
    full_ref: str
    digest: str | None = None


@dataclass(frozen=True)
class ExportedFile:
    path: Path
    arcname: str


@dataclass(frozen=True)
class DockerBuildOptions:
    build_args: dict
    platform: str | None
    no_cache: bool
    cache_from: list[str] | None = None
    target: str | None = None
    dockerfile: str | None = None


@dataclass
class DeploymentResult:
    deployed_kubernetes_manifests: list[dict] = field(default_factory=list)
    built_docker_images: list[BuiltDockerImage] = field(default_factory=list)
    exported_files: list[ExportedFile] = field(default_factory=list)


@cache
def docker_env() -> docker.DockerClient:
    return docker.from_env()


def docker_pull_for_cache(env: docker.DockerClient, full_ref: str) -> bool:
    try:
        env.images.pull(full_ref)
    except docker.errors.APIError:
        logger.debug(f'Could not pull {full_ref} for cache (not found or not accessible)')
        return False
    else:
        logger.info(f'Pulled {full_ref} for cache')
        return True


def docker_build_image(
    env: docker.DockerClient,
    context_dir: Path,
    full_ref: str,
    options: DockerBuildOptions,
) -> None:
    for line in env.api.build(
        path=str(context_dir),
        dockerfile=options.dockerfile,
        tag=full_ref,
        nocache=options.no_cache,
        pull=True,
        forcerm=True,
        buildargs=options.build_args,
        platform=options.platform,
        cache_from=options.cache_from,
        target=options.target,
        decode=True,
    ):
        logger.debug(str(line).strip())

        if 'error' in line:
            raise docker.errors.BuildError(line['error'], iter([]))


def docker_push_image(env: docker.DockerClient, repository: str, tag: str) -> str | None:
    # First try without auth_config so credential helpers (e.g. gcloud) are used.
    # Fall back to auth_config={} for anonymous pushes — needed since Docker 28.3.3
    # which rejects pushes without an X-Registry-Auth header.
    # https://github.com/docker/docker-py/issues/3348#issuecomment-3224418755
    push_kwargs: dict[str, Any] = {}
    max_attempts = 5
    backoff_base_seconds = 2
    max_backoff_seconds = 30

    # thanks gcp for failing with errors `reset reason: overflow`
    for attempt in range(max_attempts):
        error: str | None = None
        digest: str | None = None
        try:
            for line in env.images.push(
                repository=repository,
                tag=tag,
                stream=True,
                decode=True,
                **push_kwargs,
            ):
                logger.debug(str(line).strip())
                if 'error' in line:
                    error = line['error']
                    break
                if 'aux' in line:
                    digest = line['aux'].get('Digest')
        except docker.errors.DockerException as exc:
            error = str(exc)

        if error is None:
            return digest

        logger.warning(f'Push attempt {attempt + 1}/{max_attempts} failed: {error}')
        if attempt < max_attempts - 1:
            delay = min(backoff_base_seconds**attempt, max_backoff_seconds)
            logger.info(f'Retrying docker push in {delay}s')
            time.sleep(delay)

        if attempt == 0 and 'IncompleteRead' in error:
            logger.debug('Retrying push with empty auth_config (Docker 28.3.3+ workaround)')
            push_kwargs['auth_config'] = {}
            continue

        if attempt < max_attempts - 1:
            continue

        raise RuntimeError(error)
    return None


def docker_export_stage(
    env: docker.DockerClient,
    context_dir: Path,
    export: KonaChallengeConfig.ChallengeDeploymentConfig.DockerImage.Export,
    build_options: DockerBuildOptions,
    output_dir: Path,
) -> list[ExportedFile]:
    temp_tag = f'kona-export-{export.stage}:tmp'

    docker_build_image(
        env=env,
        context_dir=context_dir,
        full_ref=temp_tag,
        options=DockerBuildOptions(
            build_args=build_options.build_args,
            platform=build_options.platform,
            no_cache=build_options.no_cache,
            cache_from=build_options.cache_from,
            target=export.stage,
        ),
    )

    dst = export.dst.strip().rstrip('/')

    # Unfortunately, we can't pass buildkit args so we have to workaround like this :(
    # We are never starting this container though, so it is not as bad
    container = env.containers.create(temp_tag, command='/')
    try:
        archive_stream, _ = container.get_archive(export.src)
        raw = b''.join(archive_stream)

        # get_archive returns a tar rooted at the basename of src
        src_base = export.src.strip('/').split('/')[-1] if export.src.strip('/') else ''

        exported: list[ExportedFile] = []
        with tarfile.open(fileobj=io.BytesIO(raw), mode='r') as tar:
            for member in tar.getmembers():
                if not member.isfile():
                    continue

                rel = member.name.lstrip('./')
                if src_base and rel.startswith(f'{src_base}/'):
                    rel = rel[len(src_base) + 1 :]
                if not rel:
                    continue

                arcname = f'{dst}/{rel}' if dst else rel
                out_path = output_dir / arcname
                out_path.parent.mkdir(parents=True, exist_ok=True)

                extracted = tar.extractfile(member)
                if extracted is None:
                    continue
                out_path.write_bytes(extracted.read())

                exported.append(ExportedFile(path=out_path, arcname=arcname))

        logger.info(f'Exported {len(exported)} file(s) from stage {export.stage}')
        return exported
    finally:
        container.remove(force=True)
        with contextlib.suppress(docker.errors.APIError):
            env.images.remove(temp_tag, force=True)


async def docker_build_images(
    result: DeploymentResult,
    config: KonaGlobalConfig,
    path: Path,
    deployment_config: KonaChallengeConfig.ChallengeDeploymentConfig,
    export_dir: Path | None = None,
) -> None:
    if not deployment_config.images:
        # No need to call docker_env if we don't need it
        return

    env = docker_env()
    for image in deployment_config.images:
        repository = image.name
        if image.registry_name:
            registry = config.registries.get(image.registry_name)
            if not registry:
                msg = f'Unknown registry name "{image.registry_name}" for {image.name}:{image.tag}'
                raise ValueError(msg)

            repository = f'{registry.rstrip("/")}/{image.name}'

        full_ref = f'{repository}:{image.tag}'
        full_path = (path / image.build_context).resolve().absolute()

        logger.info(f'Building {full_ref} for {image.name}:{image.tag}')

        cache_from: list[str] | None = None
        if image.registry_name and not image.no_cache:
            pulled = await asyncio.to_thread(docker_pull_for_cache, env, full_ref)
            if pulled:
                cache_from = [full_ref]

        build_options = DockerBuildOptions(
            build_args=image.build_args,
            platform=image.platform,
            no_cache=image.no_cache,
            cache_from=cache_from,
            dockerfile=image.dockerfile,
        )

        await asyncio.to_thread(
            docker_build_image,
            env=env,
            context_dir=full_path,
            full_ref=full_ref,
            options=build_options,
        )

        digest: str | None = None
        if not image.registry_name:
            logger.warning(f'Skipping push for {full_ref} (no registry specified)')
        else:
            logger.info(f'Pushing {full_ref} for {image.name}:{image.tag}')
            digest = await asyncio.to_thread(docker_push_image, env=env, repository=repository, tag=image.tag)

        result.built_docker_images.append(
            BuiltDockerImage(
                path=full_path,
                full_ref=full_ref,
                digest=digest,
            )
        )

        if image.exports and export_dir is not None:
            for export in image.exports:
                exported = await asyncio.to_thread(
                    docker_export_stage,
                    env=env,
                    context_dir=full_path,
                    export=export,
                    build_options=build_options,
                    output_dir=export_dir,
                )
                result.exported_files.extend(exported)


def _to_dict(obj: Any) -> dict[str, Any]:  # noqa: ANN401
    return obj.to_dict() if hasattr(obj, 'to_dict') else obj


def build_manifest_context(
    config: KonaGlobalConfig,
    result: DeploymentResult,
    deployment_config: KonaChallengeConfig.ChallengeDeploymentConfig,
    challenges: list,
    *,
    use_image_digest: bool = True,
) -> dict[str, Any]:
    images: dict[str, str] = {}
    for image, built in zip(deployment_config.images, result.built_docker_images, strict=False):
        ref = built.full_ref
        if use_image_digest and built.digest:
            ref = f'{ref}@{built.digest}'
        images[image.name] = ref
    return {
        'registries': config.registries,
        'images': images,
        'config': config,
        'challenges': challenges,
    }


def _k8s_get_resource(dyn: DynamicClient, api_version: str, kind: str) -> Any:  # noqa: ANN401
    try:
        return dyn.resources.get(api_version=api_version, kind=kind)
    except ResourceNotFoundError as err:
        msg = f'Unknown resource {api_version}/{kind}. Is the CRD installed?'
        raise ValueError(msg) from err


def _k8s_apply_manifest(
    dyn: DynamicClient,
    document: dict[str, Any],
) -> tuple[str, dict[str, Any]]:
    api_version = document.get('apiVersion')
    kind = document.get('kind')
    meta = document.get('metadata') or {}
    name = meta.get('name')
    if not api_version or not kind or not name:
        msg = 'manifest must include apiVersion, kind, metadata.name'
        raise ValueError(msg)

    resource = _k8s_get_resource(dyn, api_version, kind)
    kwargs: dict[str, Any] = {}
    if getattr(resource, 'namespaced', False):
        kwargs['namespace'] = meta.get('namespace') or DEFAULT_NAMESPACE

    try:
        resource.get(name=name, **kwargs)
    except ApiException as e:
        if e.status == HTTPStatus.NOT_FOUND:
            created = resource.create(body=document, **kwargs)
            return 'created', _to_dict(created)
        raise

    # strategic merge patch
    patched = resource.patch(
        name=name,
        body=document,
        content_type='application/merge-patch+json',
        **kwargs,
    )
    return 'configured', _to_dict(patched)


def k8s_expand_manifest(doc: dict[str, Any]) -> Iterable[dict[str, Any]]:
    if doc.get('kind') == 'List' and isinstance(doc.get('items'), list):
        return (i for i in doc['items'] if isinstance(i, dict))
    return (doc,)


def _resolve_cluster_targets(config: KonaGlobalConfig, cluster_name: str | None) -> list[str]:
    clusters_count = len(config.clusters)
    if not cluster_name:
        if clusters_count != 1:
            msg = 'Cluster name should always be defined when there are more than one cluster defined'
            raise ValueError(msg)
        return []
    return resolve_cluster_names(config, cluster_name)


async def _k8s_apply_items_to_cluster(
    result: DeploymentResult,
    config: KonaGlobalConfig,
    cluster_name: str | None,
    items: list[dict[str, Any]],
) -> None:
    if cluster_name:
        load_kubeconfig(config, cluster_name)

    dyn = DynamicClient(ApiClient())

    for item in items:
        try:
            action, obj = _k8s_apply_manifest(dyn, item)
        except Exception as e:
            ctx = f'{item.get("apiVersion")}/{item.get("kind")} {item.get("metadata", {}).get("name")}'
            logger.exception(f'Failed {ctx}: {e}')
            raise

        meta = (obj or item).get('metadata', {})
        api_version = item.get('apiVersion')
        kind = item.get('kind')
        name = meta.get('name', '<unnamed>')
        ns = meta.get('namespace')
        scope = f'{ns}/' if ns else ''
        logger.info(f'k8s {api_version}/{kind}:{scope}{name} {action}')

        enriched = dict(item)
        enriched['_action'] = action
        result.deployed_kubernetes_manifests.append(enriched)


async def _k8s_apply_items(
    result: DeploymentResult,
    config: KonaGlobalConfig,
    cluster_name: str | None,
    items: list[dict[str, Any]],
) -> None:
    resolved = _resolve_cluster_targets(config, cluster_name)
    if resolved:
        for name in resolved:
            await _k8s_apply_items_to_cluster(result, config, name, items)
    else:
        await _k8s_apply_items_to_cluster(result, config, None, items)


def _inject_rollout_annotation(items: list[dict[str, Any]], annotation_path: str | None) -> None:
    if not annotation_path:
        return
    keys = annotation_path.split('.')
    timestamp = datetime.now(tz=UTC).isoformat()
    for item in items:
        target = item
        for key in keys:
            target = target.setdefault(key, {})
        target['konata.dev/deployed-at'] = timestamp


async def k8s_apply_manifests(
    result: DeploymentResult,
    config: KonaGlobalConfig,
    path: Path,
    deployment_config: KonaChallengeConfig.ChallengeDeploymentConfig,
    challenges: list,
) -> None:
    for manifest in deployment_config.kubernetes_manifests:
        manifest_context = build_manifest_context(
            config, result, deployment_config, challenges, use_image_digest=manifest.rollout_restart.image
        )
        items = [
            item
            for manifest_path_str in manifest.paths
            for doc in yaml.safe_load_all(render_template((path / manifest_path_str).read_text(), **manifest_context))
            for item in k8s_expand_manifest(doc)
        ]
        _inject_rollout_annotation(items, manifest.rollout_restart.annotation_path)
        await _k8s_apply_items(result, config, manifest.cluster_name, items)


async def k8s_apply_inline_manifests(
    result: DeploymentResult,
    config: KonaGlobalConfig,
    deployment_config: KonaChallengeConfig.ChallengeDeploymentConfig,
    challenges: list,
) -> None:
    for manifest in deployment_config.kubernetes_inline_manifests:
        manifest_context = build_manifest_context(
            config, result, deployment_config, challenges, use_image_digest=manifest.rollout_restart.image
        )
        items = [
            item
            for doc in manifest.documents
            for item in k8s_expand_manifest(render_template_values(doc, **manifest_context))
        ]
        _inject_rollout_annotation(items, manifest.rollout_restart.annotation_path)
        await _k8s_apply_items(result, config, manifest.cluster_name, items)


def _postprocess_image_names(config: KonaGlobalConfig, challenge_config: KonaChallengeConfig) -> None:
    ctx: dict[str, Any] = {'challenges': challenge_config.challenges, 'config': config}
    for image in challenge_config.deployment.images:
        image.name = render_template(image.name, **ctx)


async def deploy_challenge(
    config: KonaGlobalConfig,
    path: Path,
    challenge_config: KonaChallengeConfig,
    export_dir: Path | None = None,
) -> DeploymentResult:
    result = DeploymentResult()
    deployment_config = challenge_config.deployment
    _postprocess_image_names(config, challenge_config)

    await docker_build_images(result, config, path, deployment_config, export_dir=export_dir)
    challenges = challenge_config.challenges
    await k8s_apply_manifests(result, config, path, deployment_config, challenges)
    await k8s_apply_inline_manifests(result, config, deployment_config, challenges)
    return result
