#!/usr/bin/env python3
"""Stand-ins for Ollama (:11434) and LM Studio (:1234) so the installer's steps and the agent's
pull/status/selftest commands can be tested without downloading real runtimes or models."""
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer


class Ollama(BaseHTTPRequestHandler):
    models = {}   # name -> size
    loaded = set()

    def log_message(self, *a):
        pass

    def _json(self, obj, code=200):
        b = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def _body(self):
        n = int(self.headers.get("Content-Length") or 0)
        return json.loads(self.rfile.read(n) or b"{}")

    def do_GET(self):
        if self.path == "/api/version":
            return self._json({"version": "mock-0"})
        if self.path == "/api/tags":
            return self._json({"models": [{"name": n, "size": s} for n, s in Ollama.models.items()]})
        if self.path == "/api/ps":
            return self._json({"models": [{"name": n, "size": Ollama.models[n]} for n in Ollama.loaded]})
        self._json({"error": "nope"}, 404)

    def do_POST(self):
        b = self._body()
        if self.path == "/api/pull":
            self.send_response(200)
            self.end_headers()
            self.wfile.write((json.dumps({"status": "pulling manifest"}) + "\n").encode())
            self.wfile.write((json.dumps({"status": "pulling abc", "total": 400000000, "completed": 400000000}) + "\n").encode())
            Ollama.models[b["name"]] = 397821319
            self.wfile.write((json.dumps({"status": "success"}) + "\n").encode())
            return
        if self.path == "/api/generate":  # used to unload (keep_alive 0)
            Ollama.loaded.discard(b.get("model"))
            return self._json({})
        if self.path == "/v1/chat/completions":
            Ollama.loaded.add(b["model"])
            return self._json({"choices": [{"message": {"role": "assistant", "content": "OK"}}]})
        self._json({"error": "nope"}, 404)


class LMStudio(BaseHTTPRequestHandler):
    models = {}   # key -> size
    loaded = {}   # key -> instance id

    def log_message(self, *a):
        pass

    _json = Ollama._json
    _body = Ollama._body

    def do_GET(self):
        if self.path == "/api/v1/models":
            return self._json({"models": [
                {"type": "llm", "key": k, "size_bytes": s, "loaded_instances": ([{"id": LMStudio.loaded[k]}] if k in LMStudio.loaded else [])}
                for k, s in LMStudio.models.items()]})
        if self.path.startswith("/api/v1/models/download/status/"):
            return self._json({"status": "completed", "total_size_bytes": 400000000, "downloaded_bytes": 400000000})
        self._json({"error": "nope"}, 404)

    def do_POST(self):
        b = self._body()
        if self.path == "/api/v1/models/download":
            repo = b["model"].rstrip("/").split("/")[-1]
            key = repo.lower().replace("-gguf", "")
            LMStudio.models[key] = 397821319
            return self._json({"job_id": "job1", "status": "downloading", "total_size_bytes": 400000000})
        if self.path == "/api/v1/models/load":
            LMStudio.loaded[b["model"]] = b["model"]
            return self._json({"type": "llm", "instance_id": b["model"], "status": "loaded"})
        if self.path == "/api/v1/models/unload":
            for k, v in list(LMStudio.loaded.items()):
                if v == b["instance_id"]:
                    del LMStudio.loaded[k]
            return self._json({"instance_id": b["instance_id"]})
        if self.path == "/v1/chat/completions":
            return self._json({"choices": [{"message": {"role": "assistant", "content": "OK"}}]})
        self._json({"error": "nope"}, 404)


def serve(handler, port):
    s = ThreadingHTTPServer(("127.0.0.1", port), handler)
    threading.Thread(target=s.serve_forever, daemon=True).start()


if __name__ == "__main__":
    serve(Ollama, 11434)
    serve(LMStudio, 1234)
    threading.Event().wait()