#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Train a local operator YOLO building detector from an exported GeoIntel dataset.

Environment variables:
  OPERATOR_YOLO_DATASET_DIR   Directory containing dataset.yaml.
                               Default: /app/storage/operator-data/yolo-building-dataset
  YOLO_BASE_MODEL_PATH         Existing local base .pt model path.
                               Default: /app/models/yolov8n.pt
  TRAIN_OUTPUT_DIR             Ultralytics project output directory.
                               Default: /app/storage/training/operator-yolo
  TRAIN_RUN_NAME               Ultralytics run name.
                               Default: geointel-building-detector
  TRAIN_MODEL_OUTPUT_PATH      Destination for the best trained .pt file.
                               Default: /app/models/geointel-building-detector.pt
  TRAIN_EPOCHS                 Training epochs. Default: 8
  TRAIN_IMGSZ                  Image size. Default: 512
  TRAIN_BATCH                  Batch size. Default: 2
  TRAIN_WORKERS                Data-loader workers. Default: 0
  TRAIN_DEVICE                 Device passed to Ultralytics. Default: cpu
  PYTHON_BIN                   Python executable. Default: /opt/geointel/venv/bin/python
                               when present, otherwise python3.

This helper is an operator/runtime smoke wrapper. It requires an existing local
base model and an existing local dataset.yaml. It does not create app features.
EOF
}

if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
  usage
  exit 0
fi

OPERATOR_YOLO_DATASET_DIR="${OPERATOR_YOLO_DATASET_DIR:-/app/storage/operator-data/yolo-building-dataset}"
YOLO_BASE_MODEL_PATH="${YOLO_BASE_MODEL_PATH:-/app/models/yolov8n.pt}"
TRAIN_OUTPUT_DIR="${TRAIN_OUTPUT_DIR:-/app/storage/training/operator-yolo}"
TRAIN_RUN_NAME="${TRAIN_RUN_NAME:-geointel-building-detector}"
TRAIN_MODEL_OUTPUT_PATH="${TRAIN_MODEL_OUTPUT_PATH:-/app/models/geointel-building-detector.pt}"
TRAIN_EPOCHS="${TRAIN_EPOCHS:-8}"
TRAIN_IMGSZ="${TRAIN_IMGSZ:-512}"
TRAIN_BATCH="${TRAIN_BATCH:-2}"
TRAIN_WORKERS="${TRAIN_WORKERS:-0}"
TRAIN_DEVICE="${TRAIN_DEVICE:-cpu}"
if [[ -z "${PYTHON_BIN:-}" ]]; then
  if [[ -x "/opt/geointel/venv/bin/python" ]]; then
    PYTHON_BIN="/opt/geointel/venv/bin/python"
  else
    PYTHON_BIN="python3"
  fi
fi

DATASET_YAML="${OPERATOR_YOLO_DATASET_DIR%/}/dataset.yaml"
SUMMARY_PATH="${TRAIN_OUTPUT_DIR%/}/${TRAIN_RUN_NAME}/training_summary.json"
export DATASET_YAML
export YOLO_BASE_MODEL_PATH
export TRAIN_OUTPUT_DIR
export TRAIN_RUN_NAME
export TRAIN_MODEL_OUTPUT_PATH
export TRAIN_EPOCHS
export TRAIN_IMGSZ
export TRAIN_BATCH
export TRAIN_WORKERS
export TRAIN_DEVICE
export PYTHON_BIN
export SUMMARY_PATH

if [[ ! -f "${DATASET_YAML}" ]]; then
  echo "Dataset YAML not found: ${DATASET_YAML}" >&2
  exit 1
fi

if [[ ! -f "${YOLO_BASE_MODEL_PATH}" ]]; then
  echo "Base model file not found: ${YOLO_BASE_MODEL_PATH}" >&2
  exit 1
fi

mkdir -p "${TRAIN_OUTPUT_DIR}" "$(dirname "${TRAIN_MODEL_OUTPUT_PATH}")"

"${PYTHON_BIN}" - <<'PY'
from __future__ import annotations

import hashlib
import json
import os
import shutil
from pathlib import Path


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def seed_ultralytics_font() -> None:
    font_candidates = [
        Path("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"),
        Path("/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf"),
        Path("/usr/share/fonts/truetype/freefont/FreeSans.ttf"),
    ]
    source_font = next((font for font in font_candidates if font.exists()), None)
    if source_font is None:
        return
    config_root = Path(os.environ.get("YOLO_CONFIG_DIR", str(Path.home() / ".config")))
    target_font = config_root / "Ultralytics" / "Arial.ttf"
    target_font.parent.mkdir(parents=True, exist_ok=True)
    if not target_font.exists():
        shutil.copy2(source_font, target_font)


seed_ultralytics_font()

from ultralytics import YOLO

dataset_yaml = Path(os.environ["DATASET_YAML"])
base_model_path = Path(os.environ["YOLO_BASE_MODEL_PATH"])
train_output_dir = Path(os.environ["TRAIN_OUTPUT_DIR"])
run_name = os.environ["TRAIN_RUN_NAME"]
trained_model_output_path = Path(os.environ["TRAIN_MODEL_OUTPUT_PATH"])
summary_path = Path(os.environ["SUMMARY_PATH"])
epochs = int(os.environ["TRAIN_EPOCHS"])
image_size = int(os.environ["TRAIN_IMGSZ"])
batch_size = int(os.environ["TRAIN_BATCH"])
workers = int(os.environ["TRAIN_WORKERS"])
device = os.environ["TRAIN_DEVICE"]

model = YOLO(str(base_model_path))
model.train(
    data=str(dataset_yaml),
    epochs=epochs,
    imgsz=image_size,
    batch=batch_size,
    workers=workers,
    device=device,
    project=str(train_output_dir),
    name=run_name,
    exist_ok=True,
    pretrained=True,
    plots=False,
    verbose=True,
)

best_path = train_output_dir / run_name / "weights" / "best.pt"
if not best_path.exists():
    raise SystemExit(f"Expected trained model artifact was not created: {best_path}")

shutil.copy2(best_path, trained_model_output_path)
dataset_summary_path = next(
    (
        candidate
        for candidate in (
            dataset_yaml.parent / "yolo_tile_dataset_summary.json",
            dataset_yaml.parent / "yolo_dataset_summary.json",
        )
        if candidate.is_file()
    ),
    None,
)
summary = {
    "status": "ok",
    "dataset_yaml": str(dataset_yaml),
    "base_model_path": str(base_model_path),
    "train_output_dir": str(train_output_dir),
    "train_run_name": run_name,
    "trained_model_path": str(trained_model_output_path),
    "best_artifact_path": str(best_path),
    "dataset_yaml_sha256": sha256_file(dataset_yaml),
    "dataset_summary_path": str(dataset_summary_path) if dataset_summary_path else None,
    "dataset_summary_sha256": sha256_file(dataset_summary_path) if dataset_summary_path else None,
    "base_model_sha256": sha256_file(base_model_path),
    "trained_model_sha256": sha256_file(trained_model_output_path),
    "epochs": epochs,
    "image_size": image_size,
    "batch_size": batch_size,
    "workers": workers,
    "device": device,
}
summary_path.parent.mkdir(parents=True, exist_ok=True)
summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")
print(json.dumps(summary, indent=2, sort_keys=True))
PY