#!/usr/bin/env python3
"""Initialize, restore, or archive the organization's durable operator state."""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import sys
import tempfile

ROOT = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(ROOT), str(ROOT / "src")]


def main():
    from dotenv import load_dotenv
    from huggingface_hub import CommitOperationAdd, hf_hub_download
    from tsfm_bench.remote_state import HANDOFF_ROOTS, RemoteState, pack_files, restore_archive, utc_now
    from scripts.daemon_service import owned_process
    load_dotenv(ROOT / ".env")
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["initialize", "restore", "archive-local", "restore-handoff"])
    parser.add_argument("--repo", default=os.getenv("HF_STATE_REPO", "ThinkcatLab/LiveHouse-TS-state"))
    parser.add_argument("--root", type=Path, default=ROOT)
    args = parser.parse_args()
    root = args.root.resolve()
    pid_path = root / ".daemon.pid"
    if pid_path.exists() and owned_process(int(pid_path.read_text()), root):
        raise SystemExit("Stop this checkout's daemon before snapshotting or restoring state.")
    token = os.getenv("HF_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN is required")
    state = RemoteState(root, args.repo, token)
    if args.action == "initialize":
        manifest = state.checkpoint(initialize=True, archive_raw=True)
        print(json.dumps({"revision":state.revision, "files":len(manifest["files"]),
                          "sha256":manifest["sha256"], "checkpointed_at":manifest["checkpointed_at"]}))
    elif args.action == "restore":
        manifest = state.restore()
        print(f"Restored {len(manifest['files'])} verified files from {state.revision}")
    elif args.action == "restore-handoff":
        if any((root / path).exists() for path in HANDOFF_ROOTS):
            raise SystemExit("Use an empty --root for historical handoff restore; existing data must not be replaced.")
        info = state.api.dataset_info(args.repo)
        if not info.private:
            raise SystemExit("Handoff archives must be private")
        def download(name):
            return Path(hf_hub_download(args.repo, name, repo_type="dataset", revision=info.sha, token=token))
        pointer = json.loads(download("handoff/latest.json").read_text())
        manifest = json.loads(download(pointer["manifest"]).read_text())
        restore_archive(root, download(pointer["archive"]), manifest, allowed_roots=HANDOFF_ROOTS)
        print(f"Restored {len(manifest['files'])} verified historical files from {info.sha}")
    else:
        info = state.api.dataset_info(args.repo)
        if not info.private:
            raise SystemExit("Handoff archives must be private")
        # Include all available runtime archives, old inputs, papers and outputs.
        # Virtual environments/caches are reproducible; credentials use secrets.
        files = []
        for name in HANDOFF_ROOTS:
            files.extend(p for p in (root / name).rglob("*") if p.is_file())
        with tempfile.TemporaryDirectory(prefix="livehouse-handoff-") as temp:
            archive = Path(temp) / "operator.tar.gz"
            manifest = pack_files(root, sorted(files), archive)
            manifest.update(created_at=utc_now(), source_revision=os.getenv("LIVEHOUSE_SOURCE_REVISION", "local-migration"))
            prefix = "handoff/" + manifest["sha256"]
            commit = state.api.create_commit(args.repo, repo_type="dataset", parent_commit=info.sha,
                commit_message="Archive complete available operator data for organization handoff",
                operations=[
                    CommitOperationAdd(path_in_repo=prefix + ".tar.gz", path_or_fileobj=archive),
                    CommitOperationAdd(path_in_repo=prefix + ".json", path_or_fileobj=json.dumps(manifest).encode()),
                    CommitOperationAdd(path_in_repo="handoff/latest.json", path_or_fileobj=json.dumps({
                        "archive":prefix + ".tar.gz", "manifest":prefix + ".json",
                        "created_at":manifest["created_at"], "sha256":manifest["sha256"],
                    }).encode()),
                ])
        print(json.dumps({"revision":commit.oid, "archive":prefix + ".tar.gz", "files":len(files),
                          "bytes":manifest["bytes"], "sha256":manifest["sha256"]}))


if __name__ == "__main__":
    main()