"""Simplified module publish — source upload to platform.

The platform builds containers internally. Developers only need to:
  1. Write module code (manifest.json + backend/)
  2. Run: hla-compass publish --env dev

No Docker, no GHCR, no signing keys required.
"""

from __future__ import annotations

import io
import json
import subprocess
import sys
import time
import zipfile
from pathlib import Path

import click
from rich.console import Console
from rich.table import Table

from ..auth import Auth
from ..client import APIClient, APIError
from ..container_signing import (
    ContainerSigningError,
    default_private_key_path,
    default_public_key_path,
    generate_keypair,
    public_key_fingerprint,
    scan_image_with_trivy,
    verify_image_against_signers,
)
from ..container_signing import (
    sign_image as sign_container_image,
)
from ..signing import ModuleSigner
from ..validation import ModuleValidator
from .module import _require_supported_compute
from .utils import temporary_environment

console = Console()
DEFAULT_PLATFORM_SELECTOR = "linux/amd64"
TERMINAL_PUBLISH_STATUSES = {"completed", "approved", "failed", "cancelled", "rejected"}
_CANONICAL_COMPUTE_TYPES = {"lambda", "fargate", "batch", "sagemaker"}
_COMPUTE_TYPE_ALIASES = {"docker": "fargate", "container": "fargate"}

# Files/dirs to exclude from the source ZIP.
_EXCLUDE_PATTERNS = {
    "__pycache__",
    ".git",
    ".venv",
    "venv",
    "node_modules",
    ".mypy_cache",
    ".pytest_cache",
    ".ruff_cache",
    "dist",
    "build",
    "*.egg-info",
    ".DS_Store",
}


def _should_exclude(path: Path) -> bool:
    for part in path.parts:
        if part in _EXCLUDE_PATTERNS:
            return True
        if part.endswith(".egg-info"):
            return True
    return False


def _create_source_zip(module_dir: Path) -> bytes:
    """Create a ZIP of the module source directory."""
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for file_path in sorted(module_dir.rglob("*")):
            if not file_path.is_file():
                continue
            rel = file_path.relative_to(module_dir)
            if _should_exclude(rel):
                continue
            zf.write(file_path, str(rel))
    return buf.getvalue()


def _resolve_org_id(env: str, explicit_org_id: str | None) -> str | None:
    """Resolve the target organization ID."""
    if explicit_org_id:
        return explicit_org_id
    try:
        from ..config import Config

        config = Config()
        return config.get_active_org_id(env)
    except Exception:
        return None


def _sign_manifest_for_publish(
    manifest_for_api: dict[str, object],
    *,
    generate_keys: bool,
) -> tuple[dict[str, object], str]:
    signer = ModuleSigner()
    try:
        signer.get_public_key_string()
        signing_status = f"Keys found ({signer.get_key_fingerprint()[:16]}...)"
    except FileNotFoundError:
        if not generate_keys:
            raise
        signer.generate_keys()
        signing_status = f"Generated new keys ({signer.get_key_fingerprint()[:16]}...)"

    signed_manifest = dict(manifest_for_api)
    signed_manifest["signature"] = signer.sign_manifest(signed_manifest)
    signed_manifest["publicKey"] = signer.get_public_key_string()
    signed_manifest["keyFingerprint"] = signer.get_key_fingerprint()
    signed_manifest["signatureAlgorithm"] = signer.ALGORITHM
    signed_manifest["hashAlgorithm"] = signer.HASH_ALGORITHM
    return signed_manifest, signing_status


def _normalize_compute_type(raw: str) -> str:
    normalized = raw.strip().lower()
    normalized = _COMPUTE_TYPE_ALIASES.get(normalized, normalized)
    if normalized in _CANONICAL_COMPUTE_TYPES:
        return normalized
    raise ValueError(f"Unrecognized compute type: '{raw}'.")


def _normalize_compute_for_api(manifest: dict[str, object]) -> dict[str, object]:
    mapped = dict(manifest)
    raw = mapped.get("computeType")
    if raw is not None:
        normalized = _normalize_compute_type(str(raw))
        if normalized != raw:
            console.print(f"[dim]Mapped computeType '{raw}' -> '{normalized}' for platform compatibility.[/dim]")
        mapped["computeType"] = normalized
    return mapped


def _resolve_platform_specific_image_ref(image_ref: str, *, required: bool) -> str:
    if "@" in image_ref:
        return image_ref

    try:
        proc = subprocess.run(
            ["docker", "buildx", "imagetools", "inspect", "--format", "{{json .Manifest}}", image_ref],
            capture_output=True,
            text=True,
            check=True,
        )
        manifest = json.loads(proc.stdout or "{}")
    except (subprocess.CalledProcessError, json.JSONDecodeError) as exc:
        if required:
            raise click.ClickException(f"Failed to resolve registry digest for {image_ref}: {exc}") from exc
        return image_ref

    selected_platform = DEFAULT_PLATFORM_SELECTOR
    platform_parts = [part for part in selected_platform.split("/") if part]
    target_os = platform_parts[0] if len(platform_parts) >= 1 else "linux"
    target_arch = platform_parts[1] if len(platform_parts) >= 2 else "amd64"
    target_variant = platform_parts[2] if len(platform_parts) >= 3 else ""
    image_name = image_ref.split("@", 1)[0]
    manifests = manifest.get("manifests")
    if isinstance(manifests, list):
        for entry in manifests:
            if not isinstance(entry, dict):
                continue
            platform_info = entry.get("platform") or {}
            os_name = str(platform_info.get("os") or "").strip().lower()
            arch_name = str(platform_info.get("architecture") or "").strip().lower()
            variant_name = str(platform_info.get("variant") or "").strip().lower()
            if os_name != target_os.lower() or arch_name != target_arch.lower():
                continue
            if target_variant and variant_name != target_variant.lower():
                continue
            digest = str(entry.get("digest") or "").strip()
            if digest:
                return f"{image_name}@{digest}"

    digest = str(manifest.get("digest") or "").strip()
    if digest:
        return f"{image_name}@{digest}"
    if required:
        raise click.ClickException(f"Failed to resolve registry digest for {image_ref}")
    return image_ref


def _render_trusted_signers(trusted_signers: list[dict[str, object]]) -> Table:
    table = Table(title="Trusted Signers", show_lines=False)
    table.add_column("Type", style="cyan", no_wrap=True)
    table.add_column("Source", style="green")
    table.add_column("Details", style="magenta")
    for signer in trusted_signers:
        signer_type = str(signer.get("type") or signer.get("signer_type") or "unknown")
        details = (
            signer.get("identityRegexp")
            or signer.get("identity_regexp")
            or signer.get("publicKeyFingerprint")
            or signer.get("public_key_fingerprint")
            or signer.get("keyId")
            or signer.get("key_id")
            or ""
        )
        table.add_row(signer_type, str(signer.get("source") or ""), str(details))
    return table


def _normalized_publish_status(payload: dict[str, object]) -> str:
    status = payload.get("status")
    if status:
        return str(status).lower()
    state = payload.get("state")
    if not state:
        return ""
    state_normalized = str(state).strip().lower()
    state_map = {
        "submitted": "in_progress",
        "pending": "in_progress",
        "approved": "completed",
        "completed": "completed",
        "failed": "failed",
        "cancelled": "cancelled",
        "rejected": "rejected",
    }
    return state_map.get(state_normalized, state_normalized)


def _render_publish_status_table(status_payload: dict[str, object]) -> Table:
    table = Table(title="Publish Status", show_lines=False)
    table.add_column("Field", style="cyan", no_wrap=True)
    table.add_column("Value", style="green")
    for key in (
        "publishId",
        "status",
        "state",
        "buildStatus",
        "currentPhase",
        "failurePhase",
        "failureReason",
        "trustPolicySource",
        "startTime",
        "endTime",
        "message",
        "logsUrl",
    ):
        if key in status_payload and status_payload.get(key) is not None:
            table.add_row(key, str(status_payload.get(key)))
    return table


def _wait_for_publish_completion(
    *,
    client: APIClient,
    env: str,
    publish_id: str,
    timeout: int,
    interval: int,
) -> dict[str, object]:
    console.print("[dim]  Waiting for build to complete...[/dim]")
    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            status_resp = client.get_module_publish_status(publish_id)
            state = _normalized_publish_status(status_resp)

            if state in ("completed", "approved"):
                console.print("[green]  Module published successfully![/green]")
                module_id = status_resp.get("moduleId") or status_resp.get("module_id") or ""
                if module_id:
                    console.print(f"[dim]  Module ID: {module_id}[/dim]")
                return status_resp

            if state in ("failed", "rejected", "cancelled"):
                reason = status_resp.get("failureReason") or status_resp.get("reason") or "unknown"
                console.print(_render_publish_status_table(status_resp))
                console.print(f"[red]  Publish {state}: {reason}[/red]")
                raise click.ClickException(f"Publish {state}: {reason}")

            phase = status_resp.get("currentPhase") or ""
            if phase:
                console.print(f"[dim]  Phase: {phase}...[/dim]", end="\r")
        except APIError:
            pass

        time.sleep(interval)

    raise click.ClickException(
        f"Timed out after {timeout}s. Check status: hla-compass publish-status --env {env} {publish_id}"
    )


@click.command(name="sign-image")
@click.option("--image-ref", required=True, help="Image reference to sign")
@click.option(
    "--key-path",
    type=click.Path(path_type=Path),
    help="Cosign private key path (defaults to the SDK-managed local keypair)",
)
@click.option("--generate-keys", is_flag=True, help="Generate the default local cosign keypair if missing")
def sign_image(image_ref: str, key_path: Path | None, generate_keys: bool) -> None:
    """Sign a pre-built OCI image with the local cosign keypair."""
    try:
        resolved_key = key_path or default_private_key_path()
        if not resolved_key.exists():
            if key_path is not None or not generate_keys:
                raise FileNotFoundError(f"Cosign private key not found: {resolved_key}")
            generate_keypair()
            resolved_key = default_private_key_path()
    except (ContainerSigningError, FileExistsError, FileNotFoundError) as exc:
        raise click.ClickException(str(exc)) from exc

    signable_image_ref = _resolve_platform_specific_image_ref(image_ref, required=False)
    try:
        sign_container_image(signable_image_ref, key_path=resolved_key)
    except (ContainerSigningError, FileNotFoundError, subprocess.CalledProcessError) as exc:
        raise click.ClickException(f"Failed to sign image: {exc}") from exc

    public_key_path = default_public_key_path() if key_path is None else resolved_key.with_suffix(".pub")
    if public_key_path.exists():
        fingerprint = public_key_fingerprint(public_key_path.read_text(encoding="utf-8"))
        console.print(f"[green]✓ Signed image[/green] {signable_image_ref}")
        console.print(f"[dim]Public key:[/dim] {public_key_path}")
        console.print(f"[dim]Fingerprint:[/dim] {fingerprint}")
    else:
        console.print(f"[green]✓ Signed image[/green] {signable_image_ref}")


@click.command(name="verify-image")
@click.option("--env", required=True, type=click.Choice(["dev", "staging", "prod"]))
@click.option("--image-ref", required=True, help="Image reference to verify")
@click.option("--manifest", default="manifest.json", show_default=True, help="Manifest path")
@click.option("--scope", type=click.Choice(["org", "public"]), default="org", show_default=True)
@click.option("--org-id", help="Target organization id")
@click.option("--generate-keys", is_flag=True, help="Auto-generate manifest signing keys if missing")
def verify_image(env: str, image_ref: str, manifest: str, scope: str, org_id: str | None, generate_keys: bool) -> None:
    """Verify an image against the platform's resolved signer and scan policy."""
    manifest_path = Path(manifest)
    if not manifest_path.exists():
        raise click.ClickException(f"Manifest not found: {manifest_path}")

    with temporary_environment(env):
        auth = Auth()
        if not auth.is_authenticated():
            raise click.ClickException("Not authenticated")

        raw_manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        _require_supported_compute(raw_manifest)
        validator = ModuleValidator(manifest_path=str(manifest_path))
        validation_result = validator.run()
        if not validation_result.valid:
            console.print("[red]Manifest validation failed[/red]")
            for issue in validation_result.issues:
                console.print(f"  {issue.code}: {issue.message}")
            raise click.ClickException("Manifest validation failed")

        manifest_for_api = _normalize_compute_for_api(raw_manifest)
        try:
            manifest_for_api, signing_status = _sign_manifest_for_publish(
                manifest_for_api,
                generate_keys=generate_keys,
            )
        except FileNotFoundError as exc:
            raise click.ClickException(f"Signing keys not found: {exc}") from exc

        payload = {
            "imageRef": image_ref,
            "manifest": manifest_for_api,
            "scope": scope,
            "orgId": _resolve_org_id(env, org_id),
            "dryRun": True,
        }
        client = APIClient()
        try:
            dry_run_result = client.register_container_module(payload)
        except APIError as exc:
            raise click.ClickException(f"Backend verification failed: {exc}") from exc

        trusted_signers = dry_run_result.get("trustedSigners") or []
        if not trusted_signers:
            raise click.ClickException("Backend dry-run did not return any trusted signers")

        console.print(f"[dim]Manifest signing:[/dim] {signing_status}")
        console.print(_render_trusted_signers(trusted_signers))

        try:
            resolved_image_ref = _resolve_platform_specific_image_ref(image_ref, required=True)
            verification = verify_image_against_signers(resolved_image_ref, trusted_signers)
        except ContainerSigningError as exc:
            raise click.ClickException(str(exc)) from exc
        if not verification.get("verified"):
            raise click.ClickException(f"Image signature verification failed: {verification.get('error')}")

        console.print(
            f"[green]✓ Signature verified[/green] using {verification.get('signerType')} ({verification.get('source')})"
        )

        scan_result = scan_image_with_trivy(resolved_image_ref)
        if not scan_result.get("available"):
            raise click.ClickException("Trivy is required for verify-image but was not found in PATH")
        if scan_result.get("status") == "error":
            raise click.ClickException(str(scan_result.get("message") or "Trivy scan failed"))

        scan_table = Table(title="Security Scan", show_lines=False)
        scan_table.add_column("Field", style="cyan")
        scan_table.add_column("Value", style="green")
        scan_table.add_row("Status", str(scan_result.get("status")))
        severity_counts = scan_result.get("severityCounts") or {}
        for severity, count in severity_counts.items():
            scan_table.add_row(severity, str(count))
        console.print(scan_table)

        if scan_result.get("status") != "passed":
            failed = ", ".join(scan_result.get("failedSeverities") or [])
            raise click.ClickException(f"Image failed the platform security gate: {failed}")

        console.print("[green]✓ Image matches the current publish signer and scan policy[/green]")


@click.command()
@click.option("--env", required=True, type=click.Choice(["dev", "staging", "prod"]))
@click.option("--image-ref", help="External image reference to publish instead of source upload.")
@click.option("--org-id", help="Target organization ID. Defaults to active org for the environment.")
@click.option(
    "--scope",
    type=click.Choice(["org", "public"]),
    default="org",
    help="Module scope: 'org' (auto-approved) or 'public' (needs admin approval)",
)
@click.option("--generate-keys", is_flag=True, help="Auto-generate manifest signing keys if missing")
@click.option("--dry-run", is_flag=True, help="Validate without publishing")
@click.option("--wait", is_flag=True, help="Poll until publish completes")
@click.option("--timeout", type=int, default=600, show_default=True, help="Max seconds to wait")
@click.option("--interval", type=int, default=10, show_default=True, help="Poll interval in seconds")
@click.pass_context
def publish(ctx, env, image_ref, org_id, scope, generate_keys, dry_run, wait, timeout, interval):
    """Publish module source to the HLA-Compass platform.

    \b
    Uploads the module source code to the platform, which builds the
    container image, runs security scans, and registers the module.

    \b
    No Docker, GHCR, or signing keys required.

    \b
    Scope determines approval:
      org     — Auto-approved, only your organization can use it
      public  — Requires platform admin approval
    """
    with temporary_environment(env):
        auth = Auth()
        if not auth.is_authenticated():
            console.print("[red]Not authenticated. Run: hla-compass auth login --env {env}[/red]")
            sys.exit(1)

        manifest_path = Path("manifest.json")
        if not manifest_path.exists():
            console.print("[red]manifest.json not found. Are you in a module directory?[/red]")
            sys.exit(1)

        manifest = json.loads(manifest_path.read_text())
        module_name = manifest.get("name", "unknown")
        version = manifest.get("version", "0.0.0")

        console.print(f"[bold]Publishing {module_name} v{version}[/bold] → {env} ({scope})")
        validator = ModuleValidator(manifest_path=str(manifest_path))
        result = validator.run()
        if not result.valid:
            console.print("[red]Manifest validation failed:[/red]")
            for issue in result.issues:
                console.print(f"  {issue.code}: {issue.message}")
            sys.exit(1)
        console.print("[green]  Manifest valid[/green]")

        client = APIClient()
        selected_org_id = _resolve_org_id(env, org_id)

        if image_ref:
            _require_supported_compute(manifest)
            manifest_for_api = _normalize_compute_for_api(manifest)
            try:
                manifest_for_api, signing_status = _sign_manifest_for_publish(
                    manifest_for_api,
                    generate_keys=generate_keys,
                )
            except FileNotFoundError as exc:
                console.print(f"[red]Signing keys not found: {exc}[/red]")
                console.print("Re-run with --generate-keys.")
                sys.exit(1)

            if dry_run:
                payload = {
                    "imageRef": image_ref,
                    "manifest": manifest_for_api,
                    "scope": scope,
                    "orgId": selected_org_id,
                    "dryRun": True,
                }
                try:
                    dry_run_result = client.register_container_module(payload)
                except APIError as exc:
                    console.print(f"[red]Dry-run failed: {exc}[/red]")
                    sys.exit(1)

                console.print(f"[dim]Manifest signing:[/dim] {signing_status}")
                trusted_signers = dry_run_result.get("trustedSigners") or []
                if trusted_signers:
                    console.print(_render_trusted_signers(trusted_signers))
                console.print("[cyan]Dry run — no changes made.[/cyan]")
                return

            console.print(f"[dim]  Using external image: {image_ref}[/dim]")
            payload = {
                "imageRef": image_ref,
                "manifest": manifest_for_api,
                "scope": scope,
                "orgId": selected_org_id,
            }
            try:
                response = client.register_container_module(payload)
            except APIError as exc:
                console.print(f"[red]Publish failed: {exc}[/red]")
                sys.exit(1)

            publish_id = response.get("publishId") or response.get("id") or ""
            console.print(f"[green]  Publish accepted: {publish_id}[/green]")
            if not wait:
                console.print(f"\n[dim]Track progress: hla-compass publish-status --env {env} {publish_id}[/dim]")
                return

            _wait_for_publish_completion(
                client=client,
                env=env,
                publish_id=publish_id,
                timeout=timeout,
                interval=interval,
            )
            return

        if dry_run:
            console.print("[cyan]Dry run — no changes made.[/cyan]")
            return

        console.print("[dim]  Packaging source...[/dim]")
        source_zip = _create_source_zip(Path("."))
        size_mb = len(source_zip) / (1024 * 1024)
        console.print(f"[dim]  Source package: {size_mb:.1f} MB[/dim]")
        if size_mb > 50:
            console.print("[red]Source package exceeds 50 MB limit. Reduce the module size.[/red]")
            sys.exit(1)

        console.print("[dim]  Uploading to platform...[/dim]")
        manifest_for_api = dict(manifest)
        if manifest_for_api.get("computeType") == "docker":
            manifest_for_api["computeType"] = "fargate"

        try:
            response = client.publish_module_source(
                manifest=manifest_for_api,
                source_zip=source_zip,
                scope=scope,
                org_id=selected_org_id,
            )
        except APIError as exc:
            console.print(f"[red]Publish failed: {exc}[/red]")
            sys.exit(1)

        publish_id = response.get("publishId") or response.get("id") or ""
        console.print(f"[green]  Publish accepted: {publish_id}[/green]")
        if not wait:
            console.print(f"\n[dim]Track progress: hla-compass publish-status --env {env} {publish_id}[/dim]")
            return

        _wait_for_publish_completion(
            client=client,
            env=env,
            publish_id=publish_id,
            timeout=timeout,
            interval=interval,
        )


@click.command(name="publish-status")
@click.option("--env", required=True, type=click.Choice(["dev", "staging", "prod"]))
@click.argument("publish_id")
@click.option("--watch", is_flag=True, help="Continuously poll until terminal state")
@click.option("--interval", type=int, default=10, show_default=True, help="Poll interval seconds")
def publish_status(env, publish_id, watch, interval):
    """Check the status of a module publish job."""
    client = APIClient()

    while True:
        try:
            resp = client.get_module_publish_status(publish_id)
        except APIError as exc:
            console.print(f"[red]Error: {exc}[/red]")
            sys.exit(1)

        state = str(resp.get("status") or resp.get("state") or "unknown").upper()
        phase = resp.get("currentPhase") or ""
        console.print(f"[bold]{state}[/bold]" + (f" — {phase}" if phase else ""))

        if state in ("COMPLETED", "APPROVED", "FAILED", "REJECTED", "CANCELLED"):
            if state in ("FAILED", "REJECTED"):
                reason = resp.get("failureReason") or resp.get("reason") or ""
                if reason:
                    console.print(f"  Reason: {reason}")
            break

        if not watch:
            break

        time.sleep(interval)
