"""Secure auto-update system with Ed25519 signatures.

Only the holder of the private key can publish updates.
The public key is embedded in every client.
Updates propagate virally through the P2P network.

Security: Ed25519 (same as SSH, Signal) — computationally infeasible to forge.
"""
from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
import shutil
import sys
import time
import zipfile
from pathlib import Path
from typing import Any

from hivemind.network.protocol import Message, MsgType

log = logging.getLogger(__name__)

# ─── Ed25519 Public Key (embedded in every client) ───────────────
# This is the ONLY key that can sign valid updates.
# The private key is kept by the project maintainer.
# Replace this after generating your keypair with: python -m hivemind.network.updater genkey
PUBLISHER_PUBLIC_KEY = "0a88aea526baa9b8b2d9e0d18386019ebdf8b9cdccf2eb5f5e42a0dd61671a05"

# ─── Update manifest format ─────────────────────────────────────
# {
#   "version": "0.2.0",
#   "timestamp": 1708500000,
#   "sha256": "abc123...",       # Hash of the update zip
#   "size": 12345,
#   "changelog": "What's new...",
#   "signature": "base64..."      # Ed25519 signature of the above fields
# }


def _get_nacl():
    """Import PyNaCl (Ed25519 implementation)."""
    try:
        import nacl.signing
        import nacl.encoding
        return nacl
    except ImportError:
        return None


def generate_keypair(output_dir: Path = Path(".")) -> tuple[str, str]:
    """Generate a new Ed25519 keypair for update signing.
    
    Returns (public_key_hex, private_key_hex).
    Saves keys to files.
    """
    nacl = _get_nacl()
    if not nacl:
        raise ImportError("PyNaCl required: pip install pynacl")

    signing_key = nacl.signing.SigningKey.generate()
    verify_key = signing_key.verify_key

    private_hex = signing_key.encode(encoder=nacl.encoding.HexEncoder).decode()
    public_hex = verify_key.encode(encoder=nacl.encoding.HexEncoder).decode()

    # Save private key (KEEP SECRET!)
    priv_path = output_dir / "hivemind_update_key.private"
    priv_path.write_text(private_hex)
    priv_path.chmod(0o600)

    # Save public key
    pub_path = output_dir / "hivemind_update_key.public"
    pub_path.write_text(public_hex)

    return public_hex, private_hex


def sign_update(update_zip: Path, private_key_hex: str, version: str,
                changelog: str = "") -> dict:
    """Sign an update package. Returns the manifest dict.
    
    Args:
        update_zip: Path to the update .zip file
        private_key_hex: Hex-encoded Ed25519 private key
        version: Version string (e.g. "0.2.0")
        changelog: Human-readable changelog
    
    Returns:
        Manifest dict with signature
    """
    nacl = _get_nacl()
    if not nacl:
        raise ImportError("PyNaCl required: pip install pynacl")

    # Read and hash the update file
    data = update_zip.read_bytes()
    sha256 = hashlib.sha256(data).hexdigest()

    # Build manifest (without signature)
    manifest = {
        "version": version,
        "timestamp": int(time.time()),
        "sha256": sha256,
        "size": len(data),
        "changelog": changelog,
    }

    # Sign the canonical JSON of the manifest
    sign_data = json.dumps(manifest, sort_keys=True, ensure_ascii=False).encode()
    signing_key = nacl.signing.SigningKey(
        private_key_hex.encode(), encoder=nacl.encoding.HexEncoder
    )
    signed = signing_key.sign(sign_data)
    signature = signed.signature.hex()

    manifest["signature"] = signature
    return manifest


def sign_payload(text: str, private_key_hex: str) -> str:
    """Sign an arbitrary text payload with the Ed25519 private key.

    Returns:
        Hex-encoded signature string.
    """
    nacl = _get_nacl()
    if not nacl:
        # Fallback: HMAC-SHA256 when PyNaCl is not installed
        import hashlib, hmac as _hmac
        sig = _hmac.new(private_key_hex.encode(), text.encode(), hashlib.sha256).hexdigest()
        return "hmac:" + sig
    signing_key = nacl.signing.SigningKey(
        private_key_hex.encode(), encoder=nacl.encoding.HexEncoder
    )
    signed = signing_key.sign(text.encode())
    return signed.signature.hex()


def verify_update(manifest: dict, update_data: bytes,
                  public_key_hex: str = "") -> bool:
    """Verify an update's signature and integrity.
    
    Args:
        manifest: The update manifest dict
        update_data: Raw bytes of the update zip
        public_key_hex: Publisher's public key (uses embedded if empty)
    
    Returns:
        True if valid, False otherwise
    """
    nacl = _get_nacl()
    if not nacl:
        log.error("PyNaCl not installed — cannot verify updates")
        return False

    pub_key = public_key_hex or PUBLISHER_PUBLIC_KEY
    if not pub_key:
        log.error("No publisher public key configured")
        return False

    # Verify SHA256 of the data
    sha256 = hashlib.sha256(update_data).hexdigest()
    if sha256 != manifest.get("sha256"):
        log.error("Update hash mismatch: expected %s, got %s",
                  manifest.get("sha256"), sha256)
        return False

    # Verify Ed25519 signature
    try:
        signature = bytes.fromhex(manifest["signature"])
        # Reconstruct signed data (manifest without signature)
        verify_manifest = {k: v for k, v in manifest.items() if k != "signature"}
        sign_data = json.dumps(verify_manifest, sort_keys=True, ensure_ascii=False).encode()

        verify_key = nacl.signing.VerifyKey(
            pub_key.encode(), encoder=nacl.encoding.HexEncoder
        )
        verify_key.verify(sign_data, signature)
        log.info("Update signature verified: v%s", manifest.get("version"))
        return True

    except Exception as e:
        log.error("Update signature INVALID: %s", e)
        return False


class AutoUpdater:
    """Manages automatic updates via the P2P network.
    
    Flow:
    1. On connect, peers exchange version info
    2. If a peer has a newer version, request the update
    3. Verify signature with embedded public key
    4. Apply update and restart
    5. Inform other peers about the new version
    """

    def __init__(self, node: Any, install_dir: Path):
        self.node = node
        self.install_dir = install_dir
        self.updates_dir = install_dir / "updates"
        self.updates_dir.mkdir(exist_ok=True)
        self._current_version = node.version if hasattr(node, 'version') else "0.1.0"
        self._pending_manifest: dict | None = None
        self.pending_update: dict | None = None
        self._network = None
        # Load any pending update from disk
        pending_path = self.updates_dir / "pending.json"
        if pending_path.exists():
            try:
                self.pending_update = json.loads(pending_path.read_text())
            except Exception:
                pass

    def register_handlers(self, network):
        """Register update-related message handlers."""
        self._network = network
        network.on_message(MsgType.VERSION_INFO, self._handle_version_info)
        network.on_message(MsgType.UPDATE_AVAILABLE, self._handle_update_available)
        network.on_message(MsgType.UPDATE_REQUEST, self._handle_update_request)
        network.on_message(MsgType.UPDATE_DATA, self._handle_update_data)
        network.on_message("_version_check", self._handle_version_info)
        self._network = network

    async def _handle_version_info(self, conn, msg: Message):
        """Check if remote has a newer version."""
        remote_version = msg.payload.get("version", "0.0.0")
        if self._is_newer(remote_version, self._current_version):
            log.info("Peer %s has newer version: %s (we have %s)",
                     conn.address, remote_version, self._current_version)
            # Request update info
            await conn.send(Message(
                type=MsgType.UPDATE_REQUEST,
                sender_id=self.node.id,
                payload={"current_version": self._current_version},
            ))

    async def _handle_update_available(self, conn, msg: Message):
        """Received update announcement."""
        manifest = msg.payload.get("manifest", {})
        if self._is_newer(manifest.get("version", ""), self._current_version):
            log.info("Update available: v%s (changelog: %s)",
                     manifest.get("version"), manifest.get("changelog", ""))
            self._pending_manifest = manifest
            # Request the actual update data
            await conn.send(Message(
                type=MsgType.UPDATE_REQUEST,
                sender_id=self.node.id,
                payload={"current_version": self._current_version},
            ))

    async def _handle_update_request(self, conn, msg: Message):
        """A peer wants our update. Send it if we have it."""
        # Check if we have an update file for our current version
        update_file = self.updates_dir / f"hivemind-{self._current_version}.zip"
        manifest_file = self.updates_dir / f"hivemind-{self._current_version}.json"

        if update_file.exists() and manifest_file.exists():
            import base64
            manifest = json.loads(manifest_file.read_text())
            update_data = base64.b64encode(update_file.read_bytes()).decode()
            await conn.send(Message(
                type=MsgType.UPDATE_DATA,
                sender_id=self.node.id,
                payload={
                    "manifest": manifest,
                    "data_b64": update_data,
                },
            ))
            log.info("Sent update v%s to %s", self._current_version, conn.address)

    async def _handle_update_data(self, conn, msg: Message):
        """Received update data — verify and apply."""
        import base64

        manifest = msg.payload.get("manifest", {})
        data_b64 = msg.payload.get("data_b64", "")

        if not manifest or not data_b64:
            log.warning("Received empty update data")
            return

        version = manifest.get("version", "")
        if not self._is_newer(version, self._current_version):
            return

        log.info("Received update v%s — verifying...")
        update_data = base64.b64decode(data_b64)

        # CRITICAL: Verify signature
        if not verify_update(manifest, update_data):
            log.error("UPDATE REJECTED — invalid signature! Possible attack.")
            return

        # Save update
        update_file = self.updates_dir / f"hivemind-{version}.zip"
        update_file.write_bytes(update_data)
        manifest_file = self.updates_dir / f"hivemind-{version}.json"
        manifest_file.write_text(json.dumps(manifest, indent=2))

        log.info("Update v%s verified and saved — waiting for user confirmation", version)

        # Store as pending — user must confirm via dashboard
        self.pending_update = {
            "version": version,
            "changelog": manifest.get("changelog", ""),
            "size": len(update_data),
            "file": str(update_file),
            "sha256": manifest.get("sha256", ""),
            "received": time.time(),
        }
        # Save pending info to disk so it survives page reloads
        pending_path = self.updates_dir / "pending.json"
        pending_path.write_text(json.dumps(self.pending_update, indent=2))

    async def apply_pending(self) -> dict:
        """Apply a pending update (called from dashboard after user confirmation)."""
        if not self.pending_update:
            return {"error": "Kein Update verfuegbar"}

        version = self.pending_update["version"]
        update_file = Path(self.pending_update["file"])

        if not update_file.exists():
            self.pending_update = None
            return {"error": "Update-Datei nicht gefunden"}

        await self._apply_update(update_file, version)

        # Clean pending
        self.pending_update = None
        pending_path = self.updates_dir / "pending.json"
        if pending_path.exists():
            pending_path.unlink()

        # Announce to other peers
        if self._network:
            await self._network.broadcast(Message(
                type=MsgType.VERSION_INFO,
                sender_id=self.node.id,
                payload={"version": version},
            ))

        return {"success": True, "version": version}

    def dismiss_pending(self):
        """User dismissed the update."""
        self.pending_update = None
        pending_path = self.updates_dir / "pending.json"
        if pending_path.exists():
            pending_path.unlink()

    async def _apply_update(self, update_zip: Path, version: str):
        """Apply an update by extracting new files."""
        backup_dir = self.updates_dir / "backup"
        backup_dir.mkdir(exist_ok=True)

        try:
            # Backup current hivemind/ directory
            src = self.install_dir / "hivemind"
            if src.exists():
                backup = backup_dir / f"hivemind-{self._current_version}"
                if backup.exists():
                    shutil.rmtree(backup)
                shutil.copytree(src, backup)

            # Extract update
            with zipfile.ZipFile(update_zip) as zf:
                zf.extractall(self.install_dir)

            self._current_version = version
            log.info("Update v%s applied successfully! Restarting...", version)

            # Auto-restart
            await self._restart()

        except Exception as e:
            log.error("Update failed: %s — rolling back", e)
            # Rollback
            backup = backup_dir / f"hivemind-{self._current_version}"
            if backup.exists():
                dst = self.install_dir / "hivemind"
                if dst.exists():
                    shutil.rmtree(dst)
                shutil.copytree(backup, dst)

    async def _restart(self):
        """Restart the HiveMind process."""
        import sys
        import os

        log.info("Restarting HiveMind in 2 seconds...")
        await asyncio.sleep(2)

        try:
            # Stop the node gracefully
            if self.node:
                await self.node.stop()
        except Exception:
            pass

        # Re-exec the current process
        python = sys.executable
        args = sys.argv[:]
        log.info("Exec: %s %s", python, " ".join(args))

        if os.name == 'nt':
            # Windows: start new process, exit current
            import subprocess
            subprocess.Popen([python] + args, close_fds=True)
            os._exit(0)
        else:
            # Unix: replace current process
            os.execv(python, [python] + args)

    @staticmethod
    def _is_newer(version_a: str, version_b: str) -> bool:
        """Check if version_a is newer than version_b."""
        try:
            def parse(v):
                return tuple(int(x) for x in v.strip().split("."))
            return parse(version_a) > parse(version_b)
        except (ValueError, AttributeError):
            return False


# ─── CLI for key generation and update signing ───────────────────

def cli():
    """Command-line interface for update management."""
    import argparse

    parser = argparse.ArgumentParser(description="HiveMind Update Manager")
    sub = parser.add_subparsers(dest="command")

    # Generate keypair
    sub.add_parser("genkey", help="Generate Ed25519 keypair")

    # Sign an update
    sign_p = sub.add_parser("sign", help="Sign an update package")
    sign_p.add_argument("zipfile", help="Path to update .zip")
    sign_p.add_argument("version", help="Version string (e.g. 0.2.0)")
    sign_p.add_argument("--key", required=True, help="Path to private key file")
    sign_p.add_argument("--changelog", default="", help="Changelog text")

    # Verify
    verify_p = sub.add_parser("verify", help="Verify an update package")
    verify_p.add_argument("zipfile", help="Path to update .zip")
    verify_p.add_argument("manifest", help="Path to manifest .json")
    verify_p.add_argument("--pubkey", default="", help="Public key (hex) or path to file")

    args = parser.parse_args()

    if args.command == "genkey":
        pub, priv = generate_keypair()
        print(f"\n🔑 Schlüsselpaar generiert!\n")
        print(f"  Public Key:  {pub}")
        print(f"  Private Key: (gespeichert in hivemind_update_key.private)")
        print(f"\n  ⚠️  PRIVATE KEY SICHER AUFBEWAHREN!")
        print(f"  📋 Public Key in updater.py → PUBLISHER_PUBLIC_KEY eintragen\n")

    elif args.command == "sign":
        private_key = Path(args.key).read_text().strip()
        manifest = sign_update(Path(args.zipfile), private_key, args.version, args.changelog)
        
        out_path = Path(args.zipfile).with_suffix(".json")
        out_path.write_text(json.dumps(manifest, indent=2))
        print(f"\n✅ Update signiert: {out_path}")
        print(f"   Version: {args.version}")
        print(f"   SHA256: {manifest['sha256']}")
        print(f"   Signatur: {manifest['signature'][:32]}...\n")

    elif args.command == "verify":
        manifest = json.loads(Path(args.manifest).read_text())
        data = Path(args.zipfile).read_bytes()
        pubkey = args.pubkey
        if pubkey and Path(pubkey).exists():
            pubkey = Path(pubkey).read_text().strip()

        if verify_update(manifest, data, pubkey):
            print(f"\n✅ Signatur gültig! Update v{manifest.get('version')} ist authentisch.\n")
        else:
            print(f"\n❌ Signatur UNGÜLTIG! Update abgelehnt.\n")
            sys.exit(1)

    else:
        parser.print_help()


if __name__ == "__main__":
    cli()
