"""
solve_maze.py — jalankan model penyelesai maze di VPS (CPU saja, RAM kecil cukup).

Persiapan sekali saja (di VPS):
    pip install torch --index-url https://download.pytorch.org/whl/cpu
    pip install numpy safetensors huggingface_hub

Pemakaian:
    # unduh model dari HuggingFace lalu selesaikan maze dari file JSON
    python solve_maze.py --repo USERNAME/maze-solver --file maze.json

    # dari stdin
    echo '[[1,1,1,1],[1,2,3,1],[1,1,1,1]]' | python solve_maze.py --repo USERNAME/maze-solver

    # jalankan sebagai API HTTP kecil:  POST /solve  body: {"maze": [[...], ...]}
    python solve_maze.py --repo USERNAME/maze-solver --serve --port 8000
"""
import argparse
import json
import os
import sys


def get_model_dir(repo: str | None, local_dir: str) -> str:
    """Pakai folder lokal jika sudah ada; jika belum, unduh dari HuggingFace."""
    if os.path.exists(os.path.join(local_dir, "model.safetensors")):
        return local_dir
    if not repo:
        sys.exit(f"Model tidak ada di {local_dir}/ — berikan --repo USERNAME/NAMA_REPO")
    from huggingface_hub import snapshot_download
    print(f"Mengunduh {repo} ...", file=sys.stderr)
    return snapshot_download(repo_id=repo, local_dir=local_dir,
                             token=os.environ.get("HF_TOKEN"))  # token hanya utk repo private


def main():
    p = argparse.ArgumentParser(description="Selesaikan maze dengan MazeNet")
    p.add_argument("--repo", help="repo HuggingFace, mis. USERNAME/maze-solver")
    p.add_argument("--dir", default="maze-solver", help="folder model lokal")
    p.add_argument("--file", help="file JSON berisi maze (default: stdin)")
    p.add_argument("--no-fallback", action="store_true",
                   help="jangan pakai BFS cadangan bila model gagal")
    p.add_argument("--show", action="store_true", help="gambar maze + jalur di terminal")
    p.add_argument("--serve", action="store_true", help="jalankan API HTTP")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8000)
    a = p.parse_args()

    model_dir = get_model_dir(a.repo, a.dir)
    sys.path.insert(0, model_dir)          # maze_model.py ikut terunduh dari repo
    import torch
    from maze_model import load_model, solve, render

    torch.set_num_threads(max(1, min(2, os.cpu_count() or 1)))
    model = load_model(model_dir)
    fallback = not a.no_fallback

    if a.serve:
        from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
        import threading
        lock = threading.Lock()  # satu permintaan sekali jalan → RAM tetap kecil

        class H(BaseHTTPRequestHandler):
            def _send(self, code, obj):
                body = json.dumps(obj).encode()
                self.send_response(code)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)

            def do_GET(self):
                self._send(200, {"status": "ok"})

            def do_POST(self):
                if self.path != "/solve":
                    return self._send(404, {"error": "pakai POST /solve"})
                try:
                    n = int(self.headers.get("Content-Length", 0))
                    if n > 2_000_000:
                        return self._send(413, {"error": "maze terlalu besar"})
                    data = json.loads(self.rfile.read(n))
                    maze = data["maze"] if isinstance(data, dict) else data
                    with lock:
                        out = solve(model, maze, fallback_bfs=fallback)
                    self._send(200, out)
                except (ValueError, KeyError, TypeError) as e:
                    self._send(400, {"error": str(e)})

        print(f"API siap di http://{a.host}:{a.port}/solve", file=sys.stderr)
        ThreadingHTTPServer((a.host, a.port), H).serve_forever()
        return

    raw = open(a.file).read() if a.file else sys.stdin.read()
    import re
    data = json.loads(re.sub(r",\s*]", "]", raw))  # toleransi koma di akhir seperti "[...], ]"
    maze = data["maze"] if isinstance(data, dict) else data
    out = solve(model, maze, fallback_bfs=fallback)
    if a.show and out["path"]:
        print(render(maze, [tuple(p) for p in out["path"]]), file=sys.stderr)
    print(json.dumps(out))


if __name__ == "__main__":
    main()