#!/usr/bin/env python3
"""
Create a deployment bundle for an edge node.

This script packages everything an edge node needs to perform face recognition
for a specific section: FAISS index, student map, config, and model files.

Usage:
    python scripts/create_deploy_bundle.py --section AIML-3-A
    python scripts/create_deploy_bundle.py --section AIML-3-A --output data/deploy/
    python scripts/create_deploy_bundle.py --section AIML-3-A --include-models
"""
import argparse
import hashlib
import json
import os
import shutil
import sys
import zipfile
from datetime import datetime
from pathlib import Path

# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))


def compute_checksum(filepath: Path) -> str:
    """Compute SHA-256 checksum of a file."""
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


def create_bundle(
    section_key: str,
    output_dir: Path,
    data_dir: Path,
    models_dir: Path,
    config_dir: Path,
    include_models: bool = False,
) -> Path:
    """Create a deployment bundle zip for the given section."""

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    bundle_name = f"bundle_{section_key}_{timestamp}"
    bundle_dir = output_dir / bundle_name

    print(f"📦 Creating deployment bundle for section: {section_key}")
    print(f"   Output directory: {output_dir}")

    # Create temporary bundle directory
    bundle_dir.mkdir(parents=True, exist_ok=True)

    # ─── 1. FAISS Index ──────────────────────────────────────────────
    faiss_source = data_dir / "faiss_indices" / f"{section_key}.index"
    if faiss_source.exists():
        shutil.copy2(faiss_source, bundle_dir / "faiss_index.bin")
        print(f"   ✅ FAISS index: {faiss_source.stat().st_size / 1024:.1f} KB")
    else:
        print(f"   ⚠️  FAISS index not found: {faiss_source}")
        print(f"      Run face enrollment first for section {section_key}")
        sys.exit(1)

    # ─── 2. Student Map ──────────────────────────────────────────────
    student_map_source = data_dir / "student_maps" / f"{section_key}_map.json"
    if student_map_source.exists():
        shutil.copy2(student_map_source, bundle_dir / "student_map.json")
        with open(student_map_source) as f:
            student_map = json.load(f)
        print(f"   ✅ Student map: {len(student_map)} students")
    else:
        print(f"   ⚠️  Student map not found: {student_map_source}")
        sys.exit(1)

    # ─── 3. Edge Configuration ───────────────────────────────────────
    edge_config_source = config_dir / "edge_config.yaml"
    if edge_config_source.exists():
        shutil.copy2(edge_config_source, bundle_dir / "edge_config.yaml")
        print(f"   ✅ Edge config included")
    else:
        print(f"   ⚠️  Edge config not found: {edge_config_source}")
        # Create a minimal config
        minimal_config = f"""# SmartClass Edge Configuration
section_key: "{section_key}"
recognition:
  confidence_threshold: 0.65
  min_face_size: 80
  max_faces: 30
pipeline:
  fps_target: 15
  frame_skip: 2
metrics:
  port: 9100
  enabled: true
"""
        with open(bundle_dir / "edge_config.yaml", "w") as f:
            f.write(minimal_config)
        print(f"   ✅ Generated minimal edge config")

    # ─── 4. Model Files (optional) ───────────────────────────────────
    if include_models:
        model_files = [
            "face_detection.onnx",
            "face_recognition.onnx",
            "anti_spoof.onnx",
        ]
        models_included = 0
        model_bundle_dir = bundle_dir / "models"
        model_bundle_dir.mkdir(exist_ok=True)

        for model_file in model_files:
            model_path = models_dir / model_file
            if model_path.exists():
                shutil.copy2(model_path, model_bundle_dir / model_file)
                models_included += 1
                print(f"   ✅ Model: {model_file} ({model_path.stat().st_size / 1024 / 1024:.1f} MB)")

        if models_included == 0:
            print(f"   ⚠️  No model files found in {models_dir}")
            shutil.rmtree(model_bundle_dir)

    # ─── 5. Create Manifest ──────────────────────────────────────────
    manifest = {
        "version": "1.0",
        "section_key": section_key,
        "created_at": datetime.now().isoformat(),
        "student_count": len(student_map) if "student_map" in dir() else 0,
        "includes_models": include_models,
        "files": {},
    }

    for file_path in bundle_dir.rglob("*"):
        if file_path.is_file() and file_path.name != "manifest.json":
            rel_path = str(file_path.relative_to(bundle_dir))
            manifest["files"][rel_path] = {
                "size": file_path.stat().st_size,
                "checksum": compute_checksum(file_path),
            }

    with open(bundle_dir / "manifest.json", "w") as f:
        json.dump(manifest, f, indent=2)

    # ─── 6. Create ZIP ───────────────────────────────────────────────
    zip_path = output_dir / f"{bundle_name}.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for file_path in bundle_dir.rglob("*"):
            if file_path.is_file():
                arcname = file_path.relative_to(bundle_dir)
                zf.write(file_path, arcname)

    # Also create a "latest" symlink/copy
    latest_path = output_dir / "latest_bundle.zip"
    if latest_path.exists():
        latest_path.unlink()
    shutil.copy2(zip_path, latest_path)

    # Cleanup temporary directory
    shutil.rmtree(bundle_dir)

    bundle_checksum = compute_checksum(zip_path)
    print(f"\n✅ Bundle created successfully!")
    print(f"   Path: {zip_path}")
    print(f"   Size: {zip_path.stat().st_size / 1024:.1f} KB")
    print(f"   Checksum: {bundle_checksum}")
    print(f"   Latest: {latest_path}")

    return zip_path


def main():
    parser = argparse.ArgumentParser(
        description="Create deployment bundle for SmartClass edge node"
    )
    parser.add_argument(
        "--section", required=True, help="Section key (e.g., AIML-3-A)"
    )
    parser.add_argument(
        "--output", default="data/deploy", help="Output directory (default: data/deploy)"
    )
    parser.add_argument(
        "--data-dir", default="data", help="Data directory (default: data)"
    )
    parser.add_argument(
        "--models-dir", default="models", help="Models directory (default: models)"
    )
    parser.add_argument(
        "--config-dir", default="config", help="Config directory (default: config)"
    )
    parser.add_argument(
        "--include-models", action="store_true", help="Include model files in bundle"
    )

    args = parser.parse_args()

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    create_bundle(
        section_key=args.section,
        output_dir=output_dir,
        data_dir=Path(args.data_dir),
        models_dir=Path(args.models_dir),
        config_dir=Path(args.config_dir),
        include_models=args.include_models,
    )


if __name__ == "__main__":
    main()