#!/usr/bin/env python3
"""
Stream egitim dongusu: chunk indir -> egit -> checkpoint/HF upload -> sonraki chunk.
Takilirsa / cokurse otomatik resume.

Ornek (VM):
  export HF_TOKEN=...
  export HF_REPO=HayrettinIscan/MeshAI-Base-Models
  export HF_REPO_TYPE=model
  python3 scripts/stream_train_autoloop.py --chunk-size 8 --epochs-per-chunk 1
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
CACHE = ROOT / "data" / "stream_cache"
PROGRESS = ROOT / "data" / "stream_progress.json"
CHECKPOINT_DIR = ROOT / "checkpoints"
LOG_DIR = ROOT / "logs"
HF_PREPROCESSED = os.getenv("HF_PREPROCESSED_REPO", "HayrettinIscan/MeshAI-Preprocessed-4K")
HF_REPO = os.getenv("HF_REPO", "HayrettinIscan/MeshAI-Base-Models")
HF_REPO_TYPE = os.getenv("HF_REPO_TYPE", "model")
HF_TOKEN = os.getenv("HF_TOKEN")

REQUIRED_FILES = ("geometry_latent.npz",)
RENDER_CANDIDATES = tuple(
    f"renders/{name}"
    for i in range(4)
    for name in (f"view_{i:02d}_tex.png", f"view_{i:02d}.png")
)


def _log(msg: str) -> None:
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    line = f"[{datetime.now():%Y-%m-%d %H:%M:%S}] [stream] {msg}"
    print(line, flush=True)
    with open(LOG_DIR / "stream_train_autoloop.log", "a", encoding="utf-8") as f:
        f.write(line + "\n")


def _load_progress() -> dict:
    if PROGRESS.exists():
        try:
            return json.loads(PROGRESS.read_text(encoding="utf-8"))
        except Exception:
            pass
    return {"next_index": 0, "done_uids": [], "rounds": 0, "restarts": 0}


def _save_progress(payload: dict) -> None:
    PROGRESS.parent.mkdir(parents=True, exist_ok=True)
    PROGRESS.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def _load_manifest(token: str) -> list[dict]:
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        repo_id=HF_PREPROCESSED,
        filename="preprocessed/preprocessed_objects.json",
        repo_type="dataset",
        token=token,
    )
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    rows = [r for r in payload.get("objects", []) if r.get("status", "ready") == "ready"]
    if not rows:
        raise RuntimeError("preprocessed_objects.json bos")
    return rows


def _uid_prefix(uid: str) -> str:
    return f"preprocessed/{uid[:2]}/{uid}"


def _copy_hub_file(cached: Path, target: Path) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.exists() and target.stat().st_size > 0:
        return
    shutil.copy2(cached, target)


def _prefetch_uid(uid: str, token: str) -> bool:
    """Tek obje dosyalarini stream_cache'e indir. Eksik/404 ise False."""
    from huggingface_hub import hf_hub_download

    dest = CACHE / uid[:2] / uid
    dest.mkdir(parents=True, exist_ok=True)
    try:
        for name in REQUIRED_FILES:
            cached = Path(
                hf_hub_download(
                    repo_id=HF_PREPROCESSED,
                    filename=f"{_uid_prefix(uid)}/{name}",
                    repo_type="dataset",
                    token=token,
                )
            )
            _copy_hub_file(cached, dest / name)
        got_render = False
        for rel in RENDER_CANDIDATES:
            try:
                cached = Path(
                    hf_hub_download(
                        repo_id=HF_PREPROCESSED,
                        filename=f"{_uid_prefix(uid)}/{rel}",
                        repo_type="dataset",
                        token=token,
                    )
                )
                _copy_hub_file(cached, dest / rel)
                got_render = True
            except Exception:
                continue
        if not (dest / "geometry_latent.npz").exists():
            return False
        # Bozuk npz'yi erken yakala
        try:
            import numpy as np

            with np.load(dest / "geometry_latent.npz") as z:
                _ = z["vertex_hist"]
        except Exception as exc:
            _log(f"Bozuk latent silindi uid={uid}: {exc}")
            try:
                (dest / "geometry_latent.npz").unlink(missing_ok=True)
            except OSError:
                pass
            return False
        return True
    except Exception as exc:
        msg = str(exc).lower()
        if "429" in msg or "rate limit" in msg:
            _log(f"RATE LIMIT uid={uid}: {exc}")
            raise
        _log(f"ATLA uid={uid}: {exc}")
        return False


def _prune_cache(keep_uids: set[str]) -> None:
    if not CACHE.exists():
        return
    for xx in CACHE.iterdir():
        if not xx.is_dir():
            continue
        for uid_dir in xx.iterdir():
            if uid_dir.is_dir() and uid_dir.name not in keep_uids:
                shutil.rmtree(uid_dir, ignore_errors=True)


def _upload_checkpoint() -> None:
    latest = CHECKPOINT_DIR / "latest_model.pt"
    if not latest.exists() or not HF_TOKEN:
        return
    try:
        from huggingface_hub import HfApi

        api = HfApi(token=HF_TOKEN)
        api.upload_file(
            path_or_fileobj=str(latest),
            path_in_repo="checkpoints/latest_model.pt",
            repo_id=HF_REPO,
            repo_type=HF_REPO_TYPE,
            token=HF_TOKEN,
            commit_message="stream autoloop checkpoint",
        )
        _log(f"HF upload OK -> {HF_REPO}/checkpoints/latest_model.pt ({latest.stat().st_size // 1024} KB)")
    except Exception as exc:
        _log(f"HF upload basarisiz (devam): {exc}")


def _run_train_chunk(epochs: int, stall_sec: int, checkpoint_every: int) -> int:
    """Train subprocess; stall_sec boyunca log yoksa oldur. returncode dondur."""
    cmd = [
        sys.executable,
        "-u",
        str(ROOT / "train_pipeline.py"),
        "--mode",
        "real",
        "--epochs",
        str(epochs),
        "--data-root",
        str(CACHE),
        "--validation-every",
        "999999",
        "--checkpoint-every",
        str(checkpoint_every),
    ]
    latest = CHECKPOINT_DIR / "latest_model.pt"
    if latest.exists():
        cmd.extend(["--resume_from", str(latest)])

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    _log(f"TRAIN: {' '.join(cmd)}")

    proc = subprocess.Popen(
        cmd,
        cwd=str(ROOT),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )
    assert proc.stdout is not None
    last_out = time.time()
    while True:
        line = proc.stdout.readline()
        if line:
            print(line, end="", flush=True)
            last_out = time.time()
            if "checkpoint_saved" in line.lower():
                _upload_checkpoint()
            continue
        if proc.poll() is not None:
            break
        if time.time() - last_out > stall_sec:
            _log(f"STALL {stall_sec}s — process olduruluyor (otomatik restart)")
            proc.kill()
            try:
                proc.wait(timeout=30)
            except Exception:
                pass
            return 124
        time.sleep(0.5)
    return int(proc.returncode or 0)


def run_once(args: argparse.Namespace) -> bool:
    """Bir tam tur (manifest sonuna kadar). True=bitti, False=hata/yeniden dene."""
    if not HF_TOKEN:
        _log("HATA: HF_TOKEN yok")
        return True

    progress = _load_progress()
    rows = _load_manifest(HF_TOKEN)
    n = len(rows)
    idx = int(progress.get("next_index", 0))
    _log(f"Manifest={n} | next_index={idx} | chunk={args.chunk_size}")

    while idx < n:
        chunk = rows[idx : idx + args.chunk_size]
        uids = [str(r.get("uid") or r.get("object_id")) for r in chunk]
        _log(f"CHUNK [{idx}:{idx + len(uids)}] -> {uids[:3]}{'...' if len(uids) > 3 else ''}")

        ok_uids: list[str] = []
        for uid in uids:
            try:
                if _prefetch_uid(uid, HF_TOKEN):
                    ok_uids.append(uid)
                else:
                    progress.setdefault("skipped", []).append(uid)
            except Exception as exc:
                if "429" in str(exc) or "rate limit" in str(exc).lower():
                    _log("Rate limit — 90s bekle, sonra restart")
                    _save_progress(progress)
                    time.sleep(90)
                    return False
                progress.setdefault("skipped", []).append(uid)

        if not ok_uids:
            idx += len(chunk)
            progress["next_index"] = idx
            _save_progress(progress)
            continue

        _prune_cache(set(ok_uids))
        code = _run_train_chunk(
            epochs=args.epochs_per_chunk,
            stall_sec=args.stall_sec,
            checkpoint_every=args.checkpoint_every,
        )
        if code == 124:
            progress["restarts"] = int(progress.get("restarts", 0)) + 1
            _save_progress(progress)
            _log("Stall restart — ayni chunk tekrar denenecek")
            time.sleep(15)
            return False
        if code != 0:
            progress["restarts"] = int(progress.get("restarts", 0)) + 1
            _save_progress(progress)
            _log(f"Train exit={code} — 20s sonra restart")
            time.sleep(20)
            return False

        done = list(progress.get("done_uids", []))
        done.extend(ok_uids)
        progress["done_uids"] = done[-500:]  # dosya sismesin
        idx += len(chunk)
        progress["next_index"] = idx
        progress["rounds"] = int(progress.get("rounds", 0)) + 1
        _save_progress(progress)
        _upload_checkpoint()
        _log(f"Chunk OK | ilerleme {idx}/{n}")

    _log("Tum objeler islendi.")
    return True


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="MeshAI stream train autoloop")
    p.add_argument("--chunk-size", type=int, default=8, help="Her turda kac obje indir+egit")
    p.add_argument("--epochs-per-chunk", type=int, default=1)
    p.add_argument("--checkpoint-every", type=int, default=10)
    p.add_argument("--stall-sec", type=int, default=600, help="Log yoksa oldur (sn)")
    p.add_argument("--max-restarts", type=int, default=500)
    return p.parse_args()


def main() -> None:
    args = parse_args()
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    CACHE.mkdir(parents=True, exist_ok=True)
    _log(
        f"START chunk={args.chunk_size} epochs/chunk={args.epochs_per_chunk} "
        f"stall={args.stall_sec}s repo={HF_PREPROCESSED}"
    )
    restarts = 0
    while restarts <= args.max_restarts:
        try:
            finished = run_once(args)
            if finished:
                _log("DONE")
                return
            restarts += 1
            _log(f"Auto-restart #{restarts}")
        except KeyboardInterrupt:
            _log("Kullanici durdurdu")
            raise
        except Exception as exc:
            restarts += 1
            _log(f"Beklenmeyen hata: {exc} — restart #{restarts}")
            time.sleep(30)
    _log("max-restarts asildi, cikis")
    sys.exit(1)


if __name__ == "__main__":
    main()