import base64
import builtins
import json
import subprocess
import time
import uuid
from typing import Any, Optional, Union
from uuid import UUID

import httpx
import uuid_utils
from radicli import Arg
from wasabi import msg

from prodigy_teams_broker_sdk.models import CheckProgressRequest, CheckStartRequest
from prodigy_teams_pam_sdk import models as prodigy_teams_pam_sdk_models
from prodigy_teams_pam_sdk.models import (
    BrokerUpdating,
)

from ..cli import cli
from ..errors import (
    ConnectError,
    ConnectTimeout,
    HTTPStatusError,
    InvalidURL,
    RequestError,
)
from ..messages import Messages
from ..query import resolve_broker, resolve_broker_id, resolve_package, resolve_recipe
from ..ty import BrokerStatusCheck
from ..ui import print_info_table, print_table_with_select
from ..util import URL
from ._state import get_auth_state


def uuid7() -> uuid.UUID:
    """Generate a UUIDv7 (time-sortable) as a stdlib uuid.UUID."""
    return uuid.UUID(int=uuid_utils.uuid7().int)


DEFAULT_NAMESPACE = "prodigy-teams"
DEFAULT_RELEASE = "prodigy-teams"


def check_broker_status(
    broker: prodigy_teams_pam_sdk_models.BrokerSummary,
) -> BrokerStatusCheck:
    url = URL.parse(broker.address) / "api/v1/status"
    # Before querying the broker, verify that it has finished registration
    # TODO: also track when it is killed in PAM
    if broker.state == "creating":
        return BrokerStatusCheck.CREATING
    try:
        response = httpx.get(str(url))
        response.raise_for_status()
        payload = response.json()
        # Happy case: successfully connect to the broker and check "cluster" key
        if payload.get("cluster") == "Ready":
            return BrokerStatusCheck.RUNNING
        else:
            return BrokerStatusCheck.ISSUES
    except InvalidURL:
        return BrokerStatusCheck.INVALID_ADDRESS
    except (ConnectTimeout, ConnectError):
        # This could mean the broker is down, but could also mean we have no network access
        return BrokerStatusCheck.NOT_FOUND
    except RequestError as e:
        msg.warn(Messages.E041.format(name=broker.name), str(e))
        return BrokerStatusCheck.REQUEST_ERROR
    except HTTPStatusError as e:
        # We shouldn't really log here -- better to refactor this into the commands
        msg.warn(Messages.E041.format(name=broker.name), str(e.response.status_code))
        return BrokerStatusCheck.RESPONSE_ERROR


class BrokerSummary(prodigy_teams_pam_sdk_models.BrokerSummary):
    """
    Extends the data returned by the API with computed properties.
    This is a workaround for the models being auto-generated.
    """

    @property
    def status(self) -> str:
        status = check_broker_status(self)
        return str(status)

    @classmethod
    def get_properties(cls) -> builtins.list[str]:
        return [
            prop for prop in cls.__dict__ if isinstance(cls.__dict__[prop], property)
        ]

    def model_dump(self, **kwargs: Any) -> dict[str, Any]:  # type: ignore[override]
        include = kwargs.get("include")
        exclude = kwargs.get("exclude")
        data = super().model_dump(**kwargs)
        props = self.get_properties()
        if include:
            props = [prop for prop in props if prop in include]
        if exclude:
            props = [prop for prop in props if prop not in exclude]
        for key in props:
            data[key] = getattr(self, key)
        return data


class BrokerDetail(BrokerSummary, prodigy_teams_pam_sdk_models.BrokerDetail):
    pass


@cli.subcommand(
    "clusters",
    "list",
    # fmt: off
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(BrokerSummary.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def list(
    select: builtins.list[str] = ["id", "name", "status", "address"],
    as_json: bool = False,
) -> builtins.list[BrokerSummary]:
    """List resources on the cluster"""
    client = get_auth_state().client
    # TODO: this sequentially checks broker status which could be slow
    res = [BrokerSummary(**x.model_dump()) for x in client.broker.all(page_size=100)]
    print_table_with_select(res, select=select, as_json=as_json)
    return res


@cli.subcommand(
    "clusters",
    "info",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(BrokerDetail.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def info(
    name_or_id: Union[str, UUID],
    select: Optional[builtins.list[str]] = None,
    as_json: bool = False,
) -> BrokerDetail:
    """Get detailed info for a cluster"""
    cluster = BrokerDetail(**resolve_broker(name_or_id).model_dump())
    select = [*BrokerDetail.model_fields.keys(), *BrokerDetail.get_properties()]
    print_info_table(cluster, select=select, as_json=as_json)
    return cluster


@cli.subcommand(
    "clusters",
    "update",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
    name=Arg("--new-name", help=Messages.new_name.format(noun="cluster")),
    address=Arg("--address", help=Messages.new_address.format(noun="cluster")),
    as_json=Arg("--json", help=Messages.as_json),
)
def update(
    name_or_id: Union[str, UUID],
    name: Optional[str] = None,
    address: Optional[str] = None,
    as_json: bool = False,
) -> UUID:
    """Update the cluster info"""
    cluster = resolve_broker(name_or_id)
    # TODO: The BrokerUpdating model is currently wrong on pam, making these
    # fields required. Change the query when this is fixed.
    body = BrokerUpdating(
        id=cluster.id,
        name=name or cluster.name,
        address=address or cluster.address,
        public_key=None,
        cloud_provider=None,
        cloud=None,
        state=None,
    )
    auth = get_auth_state()
    res = auth.client.broker.update(body)
    print_info_table(res, as_json=as_json)
    return cluster.id


@cli.subcommand(
    "clusters",
    "delete",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
)
def delete(name_or_id: Union[str, UUID]) -> UUID:
    """
    Delete a cluster from PAM. This only removes PAM's record
    of it. The cluster itself will continue to exist - you need
    to shut it down separately.
    """
    auth = get_auth_state()
    client = auth.client
    cluster_id = resolve_broker_id(name_or_id)
    client.broker.delete(id=cluster_id)
    msg.good(Messages.T003.format(noun="cluster", name=cluster_id))
    return cluster_id


def _check_result(label: str, ok: bool, detail: str = "") -> tuple[str, bool]:
    """Print a check result line and return (label, ok)."""
    symbol = "\u2714" if ok else "\u2718"
    suffix = f" ({detail})" if detail else ""
    if ok:
        msg.good(f"{symbol} {label}{suffix}")
    else:
        msg.fail(f"{symbol} {label}{suffix}")
    return (label, ok)


@cli.subcommand(
    "clusters",
    "check",
    name_or_id=Arg("--cluster", help=Messages.name_or_id.format(noun="cluster")),
    s3_bucket=Arg("--s3-bucket", help=Messages.check_s3_bucket),
    nfs_path=Arg("--nfs-path", help=Messages.check_nfs_path),
    recipe_name_or_id=Arg("--recipe", help=Messages.check_recipe),
    recipe_args=Arg("--recipe-args", help=Messages.check_recipe_args),
    deep=Arg("--deep", help="Run broker-side deployment checks (K8s, NFS, DB)"),
)
def check(
    name_or_id: Optional[Union[str, UUID]] = None,
    s3_bucket: Optional[str] = None,
    nfs_path: Optional[str] = None,
    recipe_name_or_id: Optional[Union[str, UUID]] = None,
    recipe_args: Optional[str] = None,
    deep: bool = False,
) -> None:
    """Check the health of a cluster deployment.

    Runs CLI-side connectivity checks against the broker and PAM.
    Use --deep to also trigger broker-side deployment checks (K8s API,
    NFS, database, etc.).
    """
    auth = get_auth_state()
    results: builtins.list[tuple[str, bool]] = []
    # Resolve the broker to check
    if name_or_id:
        broker = resolve_broker(name_or_id)
    else:
        brokers = builtins.list(auth.client.broker.all(page_size=1).first_page().items)
        if not brokers:
            msg.fail("No clusters found. Register one with 'ptc clusters register'.")
            raise SystemExit(1)
        broker = brokers[0]
    broker_url = broker.address.rstrip("/")
    # 1. Broker API responding
    try:
        r = httpx.get(f"{broker_url}/api/v1/status", timeout=10)
        r.raise_for_status()
        payload = r.json()
        api_ok = payload.get("status") == "Ready"
        cluster_ok = payload.get("cluster") == "Ready"
        results.append(_check_result(f"Broker API responding on {broker_url}", api_ok))
        # 2. Cluster status (K8s connectivity from broker)
        if api_ok:
            results.append(
                _check_result(
                    "Cluster healthy (K8s API accessible from broker)",
                    cluster_ok,
                    detail=payload.get("cluster", "unknown"),
                )
            )
    except (httpx.ConnectError, httpx.ConnectTimeout):
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}",
                False,
                detail="connection failed",
            )
        )
    except httpx.HTTPStatusError as exc:
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}",
                False,
                detail=f"HTTP {exc.response.status_code}",
            )
        )
    except Exception as exc:
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}", False, detail=str(exc)
            )
        )
    # 3. Database connected — check via kubectl if available
    db_ok = False
    try:
        pod_result = subprocess.run(
            [
                "kubectl",
                "get",
                "pods",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-l",
                "app.kubernetes.io/name=prodigy-teams-broker",
                "-o",
                "jsonpath={.items[0].status.containerStatuses[0].ready}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if pod_result.returncode == 0 and pod_result.stdout.strip() == "true":
            db_ok = True
        results.append(_check_result("Database connected", db_ok))
    except (FileNotFoundError, subprocess.TimeoutExpired):
        # kubectl not available — skip this check
        pass
    # 4. NFS storage — check PVC is bound
    try:
        pvc_result = subprocess.run(
            [
                "kubectl",
                "get",
                "pvc",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-o",
                "jsonpath={.items[*].status.phase}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if pvc_result.returncode == 0 and pvc_result.stdout.strip():
            phases = pvc_result.stdout.strip().split()
            all_bound = all(p == "Bound" for p in phases)
            results.append(
                _check_result(
                    "NFS storage mounted",
                    all_bound,
                    detail=f"{len(phases)} PVC(s) {'all bound' if all_bound else ', '.join(phases)}",
                )
            )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    # 5. PAM sync active — check broker state in PAM
    pam_ok = broker.state not in ("creating", "error")
    results.append(_check_result("PAM sync active", pam_ok, detail=broker.state))
    # 6. Provision job — check most recent job status
    try:
        job_result = subprocess.run(
            [
                "kubectl",
                "get",
                "jobs",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-l",
                "app=provision",
                "-o",
                "jsonpath={.items[-1:].status.conditions[0].type}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if job_result.returncode == 0 and job_result.stdout.strip():
            job_status = job_result.stdout.strip()
            provision_ok = job_status == "Complete"
            results.append(
                _check_result(
                    "Provision job completed", provision_ok, detail=job_status
                )
            )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    # Deep checks: delegate to broker-side check API
    if deep:
        msg.divider("Broker-side deployment checks")
        if recipe_args:
            recipe_args_dict = json.loads(recipe_args)
            assert isinstance(recipe_args_dict, dict)
        else:
            recipe_args_dict = {}
        if recipe_name_or_id:
            recipe = resolve_recipe(recipe_name_or_id, broker_id=None)
            package = resolve_package(recipe.package_id, broker_id=None)
            recipe_name = recipe.name
            package_environment = package.container
        else:
            recipe_name = None
            package_environment = None
        start_body = CheckStartRequest(
            s3_bucket=s3_bucket,
            nfs_path=nfs_path,
            package_environment=package_environment,
            recipe_name=recipe_name,
            recipe_args=recipe_args_dict,
        )
        start_response = auth.broker_client.check.start(start_body)
        while True:
            progress_body = CheckProgressRequest(id=start_response.id)
            progress_response = auth.broker_client.check.progress(progress_body)
            if progress_response.status == "done":
                assert progress_response.report
                for check_name, result in progress_response.report.items():
                    ok = result == "success"
                    skipped = result == "skip"
                    if skipped:
                        msg.info(f"- {check_name}: skipped")
                    else:
                        results.append(_check_result(check_name, ok))
                break
            time.sleep(1)
    # Summary
    failed = [label for label, ok in results if not ok]
    if failed:
        raise SystemExit(1)


@cli.subcommand(
    "clusters",
    "nodes",
    # fmt: off
    select=Arg(
        "--select",
        help=Messages.select.format(
            opts=["name", "status", "cpu", "memory", "gpu", "pod_count"]
        ),
    ),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def nodes(
    select: builtins.list[str] = [
        "name",
        "status",
        "cpu",
        "memory",
        "gpu",
        "pod_count",
    ],
    as_json: bool = False,
) -> builtins.list[dict[str, Any]]:
    """List cluster nodes with capacity and utilization."""
    auth = get_auth_state()
    node_list = auth.broker_client.nodes.list()
    rows = [
        {
            "name": n.name,
            "status": n.status,
            "cpu": n.cpu,
            "memory": n.memory,
            "gpu": str(n.gpu) if n.gpu is not None else "-",
            "pod_count": str(n.pod_count),
        }
        for n in node_list
    ]
    print_table_with_select(rows, select=select, as_json=as_json)
    return rows


@cli.subcommand(
    "clusters",
    "rotate-key",
    name_or_id=Arg("--cluster", help=Messages.name_or_id.format(noun="cluster")),
    infra_secret=Arg(
        "--infra-secret",
        help="Name of the K8s secret holding PRODIGY_TEAMS_PRIVATE_KEY",
    ),
    namespace=Arg("--namespace", help="K8s namespace"),
)
def rotate_key(
    name_or_id: Optional[Union[str, UUID]] = None,
    infra_secret: str = "prodigy-teams-infra",
    namespace: str = "prodigy-teams",
) -> None:
    """Rotate the broker RSA keypair.

    Generates a new keypair, updates the broker record in PAM with the
    new public key, patches the K8s secret with the new private key,
    and restarts the broker deployment.
    """
    from ..key_pair import generate_key_pair

    auth = get_auth_state()

    # Resolve broker
    if name_or_id is not None:
        broker = resolve_broker(name_or_id)
    else:
        brokers = builtins.list(auth.client.broker.all(page_size=1).first_page().items)
        if not brokers:
            msg.fail("No clusters found.")
            raise SystemExit(1)
        broker = brokers[0]

    msg.info(f"Rotating key for cluster {broker.name} ({broker.id})")

    # 1. Generate new keypair
    kp = generate_key_pair()
    msg.good("Generated new RSA keypair")

    # 2. Update PAM with new public key
    auth.client.broker.update(
        BrokerUpdating(
            id=broker.id,
            name=broker.name,
            address=broker.address,
            public_key=kp.public_pem,
            cloud_provider=None,
            cloud=None,
            state=None,
        )
    )
    msg.good("Updated broker public key in PAM")

    # 3. Patch K8s secret with new private key
    private_key_b64 = base64.b64encode(kp.private_pem.encode()).decode()
    public_key_b64 = base64.b64encode(kp.public_pem.encode()).decode()
    cmd = [
        "kubectl",
        "create",
        "secret",
        "generic",
        infra_secret,
        f"--namespace={namespace}",
        f"--from-literal=PRODIGY_TEAMS_PRIVATE_KEY={private_key_b64}",
        f"--from-literal=PRODIGY_TEAMS_PUBLIC_KEY={public_key_b64}",
        "--dry-run=client",
        "-o",
        "yaml",
    ]
    create = subprocess.run(cmd, capture_output=True, text=True, check=True)
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=create.stdout,
        check=True,
        text=True,
    )
    msg.good(f"Updated K8s secret {infra_secret}")

    # 4. Restart broker deployment
    subprocess.run(
        [
            "kubectl",
            "rollout",
            "restart",
            "deployment/prodigy-teams-broker",
            f"--namespace={namespace}",
        ],
        check=True,
    )
    msg.good("Restarted broker deployment — new key will be active shortly")


# --- Cluster Settings ---


@cli.subcommand(
    "clusters",
    "settings",
    name_or_id=Arg("--cluster", help=Messages.name_or_id.format(noun="cluster")),
    as_json=Arg("--json", help=Messages.as_json),
)
def settings(
    name_or_id: Optional[Union[str, UUID]] = None,
    as_json: bool = False,
) -> None:
    """Show cluster settings (auto-update, cleanup, target versions)."""
    auth = get_auth_state()
    broker_id = resolve_broker_id(name_or_id)
    detail = auth.client.cluster_settings.ensure(broker_id)
    print_info_table(detail, as_json=as_json)


@cli.subcommand(
    "clusters",
    "settings-update",
    name_or_id=Arg("--cluster", help=Messages.name_or_id.format(noun="cluster")),
    auto_update_broker=Arg(
        "--auto-update-broker", help="Enable/disable auto-update for broker"
    ),
    auto_update_cpl=Arg("--auto-update-cpl", help="Enable/disable auto-update for CPL"),
    auto_update_recipes=Arg(
        "--auto-update-recipes", help="Enable/disable auto-update for recipes"
    ),
    cleanup_orphaned_jobs=Arg(
        "--cleanup-orphaned-jobs", help="Enable/disable cleanup of orphaned jobs"
    ),
    target_broker_version=Arg(
        "--target-broker-version", help="Set target broker version to install"
    ),
    target_cpl_version=Arg(
        "--target-cpl-version", help="Set target CPL version to install"
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def settings_update(
    name_or_id: Optional[Union[str, UUID]] = None,
    auto_update_broker: Optional[bool] = None,
    auto_update_cpl: Optional[bool] = None,
    auto_update_recipes: Optional[bool] = None,
    cleanup_orphaned_jobs: Optional[bool] = None,
    target_broker_version: Optional[str] = None,
    target_cpl_version: Optional[str] = None,
    as_json: bool = False,
) -> None:
    """Update cluster settings."""
    from prodigy_teams_pam_sdk.models import ClusterSettingsUpdating

    auth = get_auth_state()
    broker_id = resolve_broker_id(name_or_id)
    current = auth.client.cluster_settings.ensure(broker_id)
    body = ClusterSettingsUpdating(
        id=current.id,
        auto_update_broker=auto_update_broker,
        auto_update_cpl=auto_update_cpl,
        auto_update_recipes=auto_update_recipes,
        cleanup_orphaned_jobs=cleanup_orphaned_jobs,
        target_broker_version=target_broker_version,
        target_cpl_version=target_cpl_version,
    )
    updated = auth.client.cluster_settings.update(body)
    print_info_table(updated, as_json=as_json)
    msg.good("Cluster settings updated")


# --- Releases ---


@cli.subcommand(
    "clusters",
    "releases",
    as_json=Arg("--json", help=Messages.as_json),
)
def releases(as_json: bool = False) -> None:
    """List available releases."""
    auth = get_auth_state()
    items = builtins.list(auth.client.release.all(page_size=100))
    print_table_with_select(
        items,
        select=["id", "component", "version", "released_at"],
        as_json=as_json,
    )
