# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Model Packaging

Creates .familiar-model packages: weights + model card + provenance +
RSA signature. Provides verification for downloaded model packages.

Package structure (ZIP-based):
  model/           — Model weights (safetensors, GGUF, etc.)
  model_card.md    — Full provenance model card
  provenance.json  — Chain proof and metadata
  signature.json   — RSA signature over package hash

Dependencies: None (pure stdlib, cryptography optional for signing)
"""

import hashlib
import json
import logging
import os
import zipfile
from datetime import datetime
from pathlib import Path

logger = logging.getLogger(__name__)


def package_model(
    model_path: str,
    model_card: str,
    genesis_hash: str = "",
    chain_proof: dict = None,
) -> Path:
    """
    Create a .familiar-model package.

    Args:
        model_path: Path to model directory or file
        model_card: Model card Markdown content
        genesis_hash: Creator's genesis block hash
        chain_proof: Chain verification proof dict

    Returns:
        Path to the created .familiar-model package
    """
    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model not found: {model_path}")

    # Determine output path
    try:
        from familiar.core.paths import DATA_DIR

        output_dir = DATA_DIR / "training" / "packages"
    except ImportError:
        output_dir = Path.home() / ".familiar" / "data" / "training" / "packages"
    output_dir.mkdir(parents=True, exist_ok=True)

    model_name = model_path.stem if model_path.is_file() else model_path.name
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    package_path = output_dir / f"{model_name}_{timestamp}.familiar-model"

    # Build provenance metadata
    provenance = {
        "packaged_at": datetime.now().isoformat(),
        "genesis_hash": genesis_hash or _get_genesis_hash(),
        "chain_proof": chain_proof or {},
        "model_name": model_name,
    }

    # Create package
    with zipfile.ZipFile(str(package_path), "w", zipfile.ZIP_DEFLATED) as zf:
        # Add model weights
        if model_path.is_file():
            zf.write(str(model_path), f"model/{model_path.name}")
        elif model_path.is_dir():
            for root, _dirs, files in os.walk(str(model_path)):
                for fname in files:
                    fpath = Path(root) / fname
                    arcname = f"model/{fpath.relative_to(model_path)}"
                    zf.write(str(fpath), arcname)

        # Add model card
        zf.writestr("model_card.md", model_card)

        # Add provenance
        zf.writestr("provenance.json", json.dumps(provenance, indent=2, default=str))

    # Compute package hash
    with open(package_path, "rb") as f:
        package_hash = hashlib.sha256(f.read()).hexdigest()

    # Sign package
    signature = _sign_hash(package_hash)

    # Append signature (rewrite zip with signature)
    sig_data = {
        "package_hash": package_hash,
        "signature": signature,
        "genesis_hash": provenance["genesis_hash"],
        "signed_at": datetime.now().isoformat(),
    }

    # Write signature as a separate file alongside the package
    sig_path = package_path.with_suffix(".familiar-model.sig")
    with open(sig_path, "w", encoding="utf-8") as f:
        json.dump(sig_data, f, indent=2)

    logger.info("Packaged model: %s (hash: %s)", package_path.name, package_hash[:16])
    return package_path


def verify_model_package(package_path: str) -> dict:
    """
    Verify a .familiar-model package.

    Checks:
      1. Package hash integrity
      2. RSA signature (if signature file exists)
      3. Provenance metadata

    Returns dict with verification results.
    """
    path = Path(package_path)
    if not path.exists():
        return {"verified": False, "error": "Package not found"}

    result = {
        "verified": False,
        "package_path": str(path),
        "checks": {},
    }

    # Compute package hash
    with open(path, "rb") as f:
        package_hash = hashlib.sha256(f.read()).hexdigest()
    result["package_hash"] = package_hash

    # Check zip integrity
    try:
        with zipfile.ZipFile(str(path), "r") as zf:
            bad = zf.testzip()
            if bad:
                result["checks"]["zip_integrity"] = f"Corrupt file: {bad}"
                return result
            result["checks"]["zip_integrity"] = "ok"

            # Read provenance
            if "provenance.json" in zf.namelist():
                provenance = json.loads(zf.read("provenance.json"))
                result["provenance"] = provenance
                result["checks"]["provenance"] = "present"
            else:
                result["checks"]["provenance"] = "missing"

            # Check model card
            if "model_card.md" in zf.namelist():
                result["checks"]["model_card"] = "present"
            else:
                result["checks"]["model_card"] = "missing"

            # Check model files
            model_files = [n for n in zf.namelist() if n.startswith("model/")]
            result["model_files"] = len(model_files)
            result["checks"]["model_files"] = f"{len(model_files)} files"

    except zipfile.BadZipFile:
        result["checks"]["zip_integrity"] = "invalid zip"
        return result

    # Verify signature
    sig_path = path.with_suffix(".familiar-model.sig")
    if sig_path.exists():
        try:
            with open(sig_path, encoding="utf-8") as f:
                sig_data = json.load(f)

            stored_hash = sig_data.get("package_hash", "")
            if stored_hash == package_hash:
                result["checks"]["hash_match"] = "ok"
            else:
                result["checks"]["hash_match"] = "mismatch"
                return result

            # Verify RSA signature
            signature = sig_data.get("signature", "")
            if signature:
                verified = _verify_signature(package_hash, signature, sig_data.get("genesis_hash"))
                result["checks"]["signature"] = "valid" if verified else "invalid"
            else:
                result["checks"]["signature"] = "unsigned"

        except (json.JSONDecodeError, IOError):
            result["checks"]["signature"] = "error reading signature file"
    else:
        result["checks"]["signature"] = "no signature file"

    # Overall verdict
    checks = result["checks"]
    result["verified"] = (
        checks.get("zip_integrity") == "ok"
        and checks.get("provenance") == "present"
        and checks.get("hash_match", "ok") == "ok"
    )

    return result


def _get_genesis_hash() -> str:
    """Get this instance's genesis hash."""
    try:
        from familiar.skills.training.bridge import get_chain_status

        status = get_chain_status()
        return status.get("genesis_hash", "") or ""
    except Exception:
        return ""


def _sign_hash(hash_str: str) -> str:
    """Sign a hash with RSA private key. Returns hex signature or empty string."""
    try:
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.asymmetric import padding

        from familiar.core.genesis import load_or_create_keypair
        from familiar.core.paths import CHAIN_DIR

        key_dir = CHAIN_DIR / "keys"
        private_key, _ = load_or_create_keypair(key_dir)

        signature = private_key.sign(
            hash_str.encode(),
            padding.PKCS1v15(),
            hashes.SHA256(),
        )
        return signature.hex()
    except Exception:
        return ""


def _verify_signature(hash_str: str, signature_hex: str, genesis_hash: str = "") -> bool:
    """Verify RSA signature. Returns True if valid."""
    try:
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.asymmetric import padding

        from familiar.core.genesis import load_or_create_keypair
        from familiar.core.paths import CHAIN_DIR

        key_dir = CHAIN_DIR / "keys"
        _, public_key = load_or_create_keypair(key_dir)

        public_key.verify(
            bytes.fromhex(signature_hex),
            hash_str.encode(),
            padding.PKCS1v15(),
            hashes.SHA256(),
        )
        return True
    except Exception:
        return False
