#!/usr/bin/env bash
#
# setup_env.sh — Zřízení reprodukovatelného dev prostředí pro CyberBrain26
#
# Kontroluje GPU (nvidia-smi), místo na disku, vytvoří venv, nainstaluje balíčky
# s --require-hashes, stáhne modelové váhy s pinovanou revizí + SHA256 ověřením.
#
# Bezpečnost: nikdy nespouštět jako root/sudo. Síť pouze přes pip a Hugging Face.
# HF_TOKEN pochází výhradně z proměnné prostředí.

set -euo pipefail

###############################################################################
# Konfigurace
###############################################################################
DEV_ENV_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="${DEV_ENV_DIR}/.venv"
REQUIREMENTS="${DEV_ENV_DIR}/requirements.txt"
PYTHON_MIN_MAJOR=3
PYTHON_MIN_MINOR=11
MIN_DISK_GB=30
MODELS_DIR="${DEV_ENV_DIR}/models"
GENERATE_HASHES=0
SKIP_HASHES=0

usage() {
    cat <<EOF
Použití: $0 [--generate-hashes] [--models-dir <cesta>] [--help]

  --generate-hashes   Před instalací vygenerovat SHA256 hashy ze stažených wheelů
                      (vyžaduje síťový přístup k PyPI).
  --skip-hashes       Nainstalovat bez --require-hashes (použito automaticky,
                      pokud requirements.txt obsahuje placeholder hashy).
  --models-dir        Adresář pro modelové váhy (výchozí: dev-env/models).
  --help              Tato nápověda.
EOF
    exit 0
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --generate-hashes) GENERATE_HASHES=1; shift ;;
        --skip-hashes) SKIP_HASHES=1; shift ;;
        --models-dir) MODELS_DIR="$2"; shift 2 ;;
        --help) usage ;;
        *) echo "Neznámý parametr: $1" >&2; usage ;;
    esac
done

###############################################################################
# 1. Kontrola: nikdy nespouštět jako root
###############################################################################
if [[ "$(id -u)" -eq 0 ]]; then
    echo "CHYBA: Nespouštějte tento skript jako root/sudo. Spusťte jako běžný uživatel." >&2
    exit 1
fi

###############################################################################
# 2. Kontrola Pythonu
###############################################################################
python_cmd="$(command -v python3 || command -v python)"
if [[ -z "${python_cmd}" ]]; then
    echo "CHYBA: Python3 není v PATH." >&2
    exit 1
fi

python_version_full="$("${python_cmd}" --version 2>&1 | awk '{print $2}')"
python_major="$(echo "${python_version_full}" | cut -d. -f1)"
python_minor="$(echo "${python_version_full}" | cut -d. -f2)"
if [[ "${python_major}" -lt "${PYTHON_MIN_MAJOR}" ]] || \
   { [[ "${python_major}" -eq "${PYTHON_MIN_MAJOR}" ]] && [[ "${python_minor}" -lt "${PYTHON_MIN_MINOR}" ]]; }; then
    echo "CHYBA: Vyžadován Python ${PYTHON_MIN_MAJOR}.${PYTHON_MIN_MINOR}+, nainstalován ${python_version_full}." >&2
    exit 1
fi
echo "INFO: Python ${python_version_full} OK."

###############################################################################
# 3. Kontrola GPU (nvidia-smi) — informativní, není povinné
###############################################################################
if command -v nvidia-smi >/dev/null 2>&1; then
    echo "INFO: NVIDIA GPU detekováno (nvidia-smi dostupný):"
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true
else
    echo "INFO: nvidia-smi nenalezen — poběží se v CPU režimu. Wan2.1 a Dia vyžadují GPU."
fi

###############################################################################
# 4. Kontrola místa na disku (df, ekvivalent shutil.disk_usage)
###############################################################################
available_kb="$(df -P "${DEV_ENV_DIR}" | awk 'NR==2 {print $4}')"
available_gb=$(( available_kb / 1024 / 1024 ))
if [[ "${available_gb}" -lt "${MIN_DISK_GB}" ]]; then
    echo "CHYBA: Nedostatek místa na disku. K dispozici ~${available_gb} GB, vyžadováno ~${MIN_DISK_GB} GB." >&2
    exit 1
fi
echo "INFO: Místo na disku OK (~${available_gb} GB volno)."

###############################################################################
# 5. Vytvoření venv
###############################################################################
if [[ ! -x "${VENV_DIR}/bin/python" ]]; then
    echo "INFO: Vytvářím venv v ${VENV_DIR}..."
    "${python_cmd}" -m venv "${VENV_DIR}"
else
    echo "INFO: venv již existuje v ${VENV_DIR}."
fi

export PATH="${VENV_DIR}/bin:$PATH"
VENV_PYTHON="${VENV_DIR}/bin/python"

"${VENV_PYTHON}" -m pip install --upgrade pip

###############################################################################
# 6. Vygenerování SHA256 hashů (volitelné, před instalací)
###############################################################################
if [[ "${GENERATE_HASHES}" -eq 1 ]]; then
    echo "INFO: Generuji SHA256 hashy pro requirements.txt..."
    tmp_wheelhouse="$(mktemp -d)"
    trap 'rm -rf "${tmp_wheelhouse}"' EXIT

    while IFS= read -r line; do
        [[ "${line}" =~ ^[[:space:]]*# ]] && continue
        [[ -z "${line// }" ]] && continue

        pkg_spec="$(echo "${line}" | awk '{print $1}')"
        if [[ -z "${pkg_spec}" ]]; then continue; fi

        echo "  Stahuji ${pkg_spec} pro výpočet hashů..."
        "${VENV_PYTHON}" -m pip download "${pkg_spec}" -d "${tmp_wheelhouse}" --no-deps || {
            echo "VAROVÁNÍ: Nelze stáhnout ${pkg_spec}, hash bude nutné doplnit ručně." >&2
            continue
        }
    done < "${REQUIREMENTS}"

    echo "INFO: Hashy jsou připraveny — spusťte 'pip hash' na stažených wheelích a aktualizujte requirements.txt."
fi

###############################################################################
# 7. Instalace balíčků (s --require-hashes, pokud hashy jsou skutečné)
###############################################################################
USE_REQUIRE_HASHES=0
if [[ "${SKIP_HASHES}" -eq 1 ]]; then
    echo "INFO: --skip-hashes — instalace bez ověřování hashů."
elif grep -q "PLACEHOLDER_REPLACE_ME_64_HEX_CHARS" "${REQUIREMENTS}" 2>/dev/null; then
    echo "VAROVÁNÍ: requirements.txt obsahuje placeholder SHA256 hashy — instalace bez --require-hashes."
else
    USE_REQUIRE_HASHES=1
fi

if [[ "${USE_REQUIRE_HASHES}" -eq 1 ]]; then
    echo "INFO: Instaluji balíčky s ověřováním SHA256 hashů (--require-hashes)..."
    "${VENV_PYTHON}" -m pip install --require-hashes -r "${REQUIREMENTS}"
else
    echo "INFO: Instaluji balíčky bez ověřování hashů."
    "${VENV_PYTHON}" -m pip install -r "${REQUIREMENTS}"
fi

echo "INFO: Balíčky nainstalovány."

###############################################################################
# 8. Stažení modelových vah (revision pinning + SHA256 ověření, bez trust_remote_code)
###############################################################################
mkdir -p "${MODELS_DIR}"

export DEV_ENV_MODELS_DIR="${MODELS_DIR}"

"${VENV_PYTHON}" - <<'PYEOF'
import hashlib
import os
import sys
from pathlib import Path

from huggingface_hub import hf_hub_download

MODELS_DIR = Path(os.environ.get("DEV_ENV_MODELS_DIR", "models"))
HF_TOKEN = os.environ.get("HF_TOKEN")  # výhradně z env, žádné výchozí hodnoty

# (repo_id, soubor, revision_sha, očekávaný_sha256 nebo None pokud se má teprve vypočítat)
MODEL_FILES = [
    # Wan2.1 1.3B text-to-video (Apache 2.0, ~8 GB VRAM)
    # ("Wan-AI/Wan2.1-T2V-1.3B", "diffusion_pytorch_model.safetensors", "<revision_sha>", None),
    # Kokoro-82M TTS (Apache 2.0, ~312 MB, CPU feasible)
    # ("hexgrad/Kokoro-82M", "kokoro-v0_19.pth", "<revision_sha>", None),
    # Dia-1.6B (Apache 2.0, ~8 GB VRAM GPU minimum)
    # ("narisai/dia-1.6B", "model.safetensors", "<revision_sha>", None),
]

def verify_sha256(filepath: Path, expected: str | None) -> str:
    sha = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            sha.update(chunk)
    digest = sha.hexdigest()
    if expected and digest != expected:
        print(f"CHYBA: SHA256 neshoda pro {filepath.name}: očekáváno {expected}, erhalteno {digest}", file=sys.stderr)
        sys.exit(1)
    return digest

if not MODEL_FILES:
    print("INFO: Žádné modelové soubory ke stažení (MODEL_FILES je prázdné — doplňte revize a hashy v setup_env.sh).")
    sys.exit(0)

for repo_id, filename, revision, expected_hash in MODEL_FILES:
    if not revision:
        print(f"CHYBA: Revize není pinována pro {repo_id}/{filename} — doplňte <revision_sha>.", file=sys.stderr)
        sys.exit(1)

    print(f"INFO: Stahuji {repo_id}/{filename} (revision={revision})...")
    downloaded = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        revision=revision,
        token=HF_TOKEN,
        local_dir=str(MODELS_DIR),
        local_dir_use_symlinks=False,
        trust_remote_code=False,  # NIKDY True
    )
    actual_hash = verify_sha256(Path(downloaded), expected_hash)
    print(f"  SHA256: {actual_hash}  OK")

print("INFO: Všechny modelové váhy staženy a ověřeny.")
PYEOF

echo "INFO: Setup dokončen. Aktivujte venv: source ${VENV_DIR}/bin/activate"