#!/usr/bin/env python3
"""Mock llama-server for exercising zshell's LLM error handling.

Implements the OpenAI-compatible subset of llama-server that the C++
backend talks to, and lets you trigger specific failure modes from the
QML chat UI by typing a keyword as your message.

Wire format is taken from ggml-org/llama.cpp (tools/server):
  - SSE events are "data: <json>\n\n" (note the space after "data:").
  - A clean stream ends with "data: [DONE]\n\n".
  - Errors are {"error": {"code", "message", "type"}}.
  - A stream that dies mid-way sends the error event, then closes
    WITHOUT [DONE].
  - A request that errors before the first token is a plain HTTP 4xx/5xx
    with the error JSON body (not a stream).
  - /props is a JSON object; n_ctx lives under
    default_generation_settings.n_ctx.

Run:
    python3 scripts/mock-llama-server.py [--port 8080]

Then point zshell's LLM endpoint at http://localhost:8080.

Trigger words (type as the chat message, case-insensitive):
    normal            -> healthy stream, ends with [DONE]  (also the default)
    instant error     -> 500 before any token ("server_error")
    context overflow  -> 400 "exceeds the available context size (2048 ...)";
                          backend retries once at 2048, then fails
    error during      -> streams a few tokens, then an SSE error event,
                          closes without [DONE]
    no done           -> streams a few tokens, closes without [DONE] and
                          without an error (server killed mid-response)
    200 error         -> HTTP 200 with a plain JSON error body (not SSE)
    tool call         -> streams reasoning, then a tool_calls round calling
                          readfile on /etc/os-release (finish_reason
                          "tool_calls"); after you approve (or deny) the
                          permission dialog, the follow-up round that
                          carries the tool result gets a normal summary
                          answer. Exercises the tool permission dialog,
                          tool execution and the multi-round transcript.
    multiple tool     -> two tool calls in one round (readfile + webfetch
    calls                 of this mock's own / endpoint); approve both at
                          once
    unknown tool      -> a tool call for a tool the client does not have;
                          approving it yields an "unknown tool" error
                          result

Control endpoint (for the non-chat routes / global state), via curl:
    curl -XPOST localhost:8080/_mock -d '{"loading": true}'
        /props, /v1/models, /health all answer 503 "Loading model"
    curl -XPOST localhost:8080/_mock -d '{"no_models": true}'
        /v1/models returns an empty list
    curl -XPOST localhost:8080/_mock -d '{"chat_mode": "error_during"}'
        force a chat behaviour regardless of the message text
        (values: normal|instant_error|context_overflow|error_during|
                 no_done|200_error|tool_call|tool_call_multi|
                 tool_call_unknown, or null to go back to keywords)
    curl -XPOST localhost:8080/_mock -d '{"reset": true}'
"""

import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

MODEL_ID = "mock-model"
CTX_SIZE = 8192

# Set in main(); used by the tool-call behaviours so webfetch can point
# at this server itself.
_host = "127.0.0.1"
_port = 8080

# Mutable, shared across handler threads.
_state_lock = threading.Lock()
_state = {
	"loading": False,  # 503 "Loading model" on /props, /v1/models, /health
	"no_models": False,  # /v1/models returns an empty list
	"chat_mode": None,  # force a chat behaviour, overriding keywords
}


def _now():
	return int(time.time())


def _chunk(cid, delta, finish_reason=None):
	return {
		"id": cid,
		"object": "chat.completion.chunk",
		"created": _now(),
		"model": MODEL_ID,
		"system_fingerprint": "mock",
		"choices": [
			{
				"index": 0,
				"delta": delta,
				"finish_reason": finish_reason,
			}
		],
	}


def _usage_chunk(cid, prompt_tokens, completion_tokens):
	return {
		"id": cid,
		"object": "chat.completion.chunk",
		"created": _now(),
		"model": MODEL_ID,
		"choices": [],
		"usage": {
			"prompt_tokens": prompt_tokens,
			"completion_tokens": completion_tokens,
			"total_tokens": prompt_tokens + completion_tokens,
		},
	}


def _error(code, message, etype):
	return {"error": {"code": code, "message": message, "type": etype}}


def _sse(obj):
	# Exact llama.cpp format_oai_sse shape: "data: " + compact json + "\n\n".
	return ("data: " + json.dumps(obj, separators=(",", ":")) + "\n\n").encode()


def _sse_done():
	return b"data: [DONE]\n\n"


# Chat behaviours. Each returns ("json", status, body_bytes) for a
# non-streaming answer, or ("sse", 200, None) and then writes to the
# provided stream writer.
def _behaviour_for(text):
	t = text.lower()
	if "instant error" in t:
		return "instant_error"
	if "context overflow" in t:
		return "context_overflow"
	if "error during" in t:
		return "error_during"
	if "no done" in t or "truncate" in t:
		return "no_done"
	if "200 error" in t:
		return "200_error"
	# Order matters: "multiple tool calls" contains "tool call".
	if "two tool calls" in t or "two tool call" in t:
		return "tool_call_duo"
	if "some tool calls" in t:
		return "tool_call_some"
	if "many tool calls" in t:
		return "tool_call_many"
	if "unknown tool" in t:
		return "tool_call_unknown"
	if "tool call" in t:
		return "tool_call"
	return "normal"


_TOOL_MODES = (
	"tool_call",
	"tool_call_duo",
	"tool_call_some",
	"tool_call_many",
	"tool_call_unknown",
)


def _tool_call_delta(index, call_id=None, name=None, arguments=None):
	call = {"index": index}
	if call_id is not None:
		call["id"] = call_id
	if name is not None or arguments is not None:
		function = {}
		if name is not None:
			function["name"] = name
		if arguments is not None:
			function["arguments"] = arguments
		call["function"] = function
	return {"tool_calls": [call]}


def _stream_tool_calls(w, cid, reasoning, calls):
	# One tool round: optional reasoning, then the calls (llama.cpp
	# streams the id/name with the first fragment of each call and the
	# arguments as successive fragments), then finish_reason
	# "tool_calls". The client concatenates the argument fragments and
	# shows the permission dialog on stream end.
	w(_sse(_chunk(cid, {"role": "assistant"})))
	if reasoning:
		for piece in reasoning:
			w(_sse(_chunk(cid, {"reasoning_content": piece})))
			time.sleep(0.05)
	for i, (call_id, name, arguments) in enumerate(calls):
		w(_sse(_chunk(cid, _tool_call_delta(i, call_id=call_id, name=name))))
		time.sleep(0.05)
		for fragment in arguments:
			w(_sse(_chunk(cid, _tool_call_delta(i, arguments=fragment))))
			time.sleep(0.05)
	w(_sse(_chunk(cid, {}, finish_reason="tool_calls")))
	w(_sse(_usage_chunk(cid, 24, len(reasoning or []) + len(calls) * 2)))
	w(_sse_done())


def _tool_result_snippets(messages, limit=160):
	snippets = []
	for msg in messages or []:
		if msg.get("role") != "tool":
			continue
		content = msg.get("content")
		if isinstance(content, list):
			content = " ".join(
				part.get("text", "")
				for part in content
				if isinstance(part, dict) and part.get("type") == "text"
			)
		flat = " ".join(str(content or "").split())
		snippets.append(flat[:limit] + ("..." if len(flat) > limit else ""))
	return snippets


def _stream_tool_round_done(w, cid, messages):
	# Follow-up round: the client sent the tool results back; answer
	# with a plain content summary so the turn completes.
	w(_sse(_chunk(cid, {"role": "assistant"})))
	parts = ["mock:", "tool", "round", "complete."]
	for i, snippet in enumerate(_tool_result_snippets(messages), 1):
		parts.append(f"result {i}:")
		parts.append(snippet or "(empty)")
	for word in parts:
		w(_sse(_chunk(cid, {"content": word + " "})))
		time.sleep(0.02)
	w(_sse(_chunk(cid, {}, finish_reason="stop")))
	w(_sse(_usage_chunk(cid, 40, len(parts) + 1)))
	w(_sse_done())


def _stream_normal(w, cid, words):
	w(_sse(_chunk(cid, {"role": "assistant"})))
	for word in words:
		w(_sse(_chunk(cid, {"content": word + " "})))
		time.sleep(0.05)
	w(_sse(_chunk(cid, {}, finish_reason="stop")))
	w(_sse(_usage_chunk(cid, 12, len(words) + 1)))
	w(_sse_done())


def _stream_error_during(w, cid, words):
	w(_sse(_chunk(cid, {"role": "assistant"})))
	for word in words:
		w(_sse(_chunk(cid, {"content": word + " "})))
		time.sleep(0.05)
	# Mid-stream failure: error event, then the connection closes without
	# [DONE] (llama.cpp terminates the stream on a task error).
	w(_sse(_error(500, "Mock: generation failed mid-stream", "server_error")))


def _stream_no_done(w, cid, words):
	w(_sse(_chunk(cid, {"role": "assistant"})))
	for word in words:
		w(_sse(_chunk(cid, {"content": word + " "})))
		time.sleep(0.05)
	# Clean FIN, no [DONE], no error: server died mid-response.


def _last_user_text(messages):
	for msg in reversed(messages or []):
		if msg.get("role") == "user":
			content = msg.get("content")
			if isinstance(content, list):
				# multimodal shape: pull out the text parts
				parts = []
				for part in content:
					if isinstance(part, dict) and part.get("type") == "text":
						parts.append(part.get("text", ""))
				content = " ".join(parts)
			return content or ""
	return ""


class Handler(BaseHTTPRequestHandler):
	protocol_version = "HTTP/1.1"
	server_version = "MockLlama/1.0"

	# --- low-level response helpers -----------------------------------
	def _send_json(self, status, obj):
		body = json.dumps(obj, separators=(",", ":")).encode()
		self.send_response(status)
		self.send_header("Content-Type", "application/json")
		self.send_header("Content-Length", str(len(body)))
		self.send_header("Connection", "close")
		self.end_headers()
		self.wfile.write(body)

	def _send_loading_error(self):
		# llama.cpp middleware: server not ready -> 503 on every route.
		self._send_json(503, _error(503, "Loading model", "unavailable_error"))

	def _not_ready(self):
		with _state_lock:
			return _state["loading"]

	def _stream(self, status, writer):
		self.send_response(status)
		self.send_header("Content-Type", "text/event-stream")
		self.send_header("Cache-Control", "no-cache")
		self.send_header("Connection", "close")
		self.end_headers()
		try:
			writer(self.wfile.write)
		except (BrokenPipeError, ConnectionResetError):
			pass
		finally:
			try:
				self.wfile.flush()
			except (BrokenPipeError, ConnectionResetError):
				pass

	# --- routing -------------------------------------------------------
	def do_GET(self):
		path = self.path.split("?", 1)[0].rstrip("/") or "/"
		if path in ("/health", "/v1/health"):
			if self._not_ready():
				return self._send_loading_error()
			return self._send_json(200, {"status": "ok"})
		if path in ("/props", "/v1/props"):
			if self._not_ready():
				return self._send_loading_error()
			return self._send_json(
				200,
				{
					"total_slots": 1,
					"model_alias": MODEL_ID,
					"model_path": "/mock/model.gguf",
					"default_generation_settings": {
						"params": {"n_predict": 2048},
						"n_ctx": CTX_SIZE,
					},
					"is_sleeping": False,
				},
			)
		if path in ("/models", "/v1/models"):
			if self._not_ready():
				return self._send_loading_error()
			with _state_lock:
				no_models = _state["no_models"]
			data = (
				[]
				if no_models
				else [
					{
						"id": MODEL_ID,
						"object": "model",
						"created": _now(),
						"owned_by": "mock",
					}
				]
			)
			return self._send_json(200, {"object": "list", "data": data})
		if path == "/":
			return self._send_json(200, {"status": "ok", "mock": True})
		return self._send_json(404, _error(404, "not found", "not_found_error"))

	def do_POST(self):
		path = self.path.split("?", 1)[0].rstrip("/") or "/"
		length = int(self.headers.get("Content-Length") or 0)
		raw = self.rfile.read(length) if length else b""
		try:
			body = json.loads(raw) if raw else {}
		except json.JSONDecodeError:
			body = {}

		if path == "/_mock":
			return self._mock_control(body)

		if path in ("/chat/completions", "/v1/chat/completions"):
			if self._not_ready():
				return self._send_loading_error()
			return self._chat(body)

		return self._send_json(404, _error(404, "not found", "not_found_error"))

	# --- /_mock control ------------------------------------------------
	def _mock_control(self, body):
		with _state_lock:
			if body.get("reset"):
				_state["loading"] = False
				_state["no_models"] = False
				_state["chat_mode"] = None
			if "loading" in body:
				_state["loading"] = bool(body["loading"])
			if "no_models" in body:
				_state["no_models"] = bool(body["no_models"])
			if "chat_mode" in body:
				_state["chat_mode"] = body["chat_mode"]
			snapshot = dict(_state)
		self.log_message("control -> %s", json.dumps(snapshot))
		return self._send_json(200, {"ok": True, "state": snapshot})

	# --- /chat/completions --------------------------------------------
	def _chat(self, body):
		text = _last_user_text(body.get("messages"))
		with _state_lock:
			forced = _state["chat_mode"]
		mode = forced or _behaviour_for(text)
		stream = bool(body.get("stream", True))
		cid = "chatcmpl-mock"
		words = ["mock", "says", "hello", "from", "the", "mock", "server"]
		self.log_message("chat mode=%s stream=%s text=%r", mode, stream, text)

		if mode == "instant_error":
			# Error before the first token: plain HTTP error, not a stream.
			return self._send_json(
				500, _error(500, "Mock: instant server error", "server_error")
			)

		if mode == "context_overflow":
			return self._send_json(
				400,
				_error(
					400,
					"request (100000 tokens) exceeds the available context "
					"size (2048 tokens), try increasing it",
					"exceed_context_size_error",
				),
			)

		if mode == "200_error":
			# 200 with a plain JSON error body (not SSE) -- the case the
			# data: filter would otherwise swallow.
			return self._send_json(
				200, _error(400, "Mock: 200-with-json-error", "server_error")
			)

		if mode in _TOOL_MODES:
			# Continuation round: the previous round's tool results came
			# back through the client (role: "tool" messages).
			if any(m.get("role") == "tool" for m in body.get("messages") or []):
				return self._stream(
					200,
					lambda w: _stream_tool_round_done(
						w, cid, body.get("messages")
					),
				)
			if mode == "tool_call":
				calls = [
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					)
				]
				reasoning = [
					"The user wants a tool call. ",
					"I will read /etc/os-release.",
				]
			elif mode == "tool_call_duo":
				calls = [
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
				]
				reasoning = [
					"Two independent tool calls, ",
					"batched together.",
				]
			elif mode == "tool_call_some":
				calls = [
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
				]
				reasoning = [
					"Two independent tool calls, ",
					"batched together.",
				]
			elif mode == "tool_call_many":
				calls = [
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							'{"url": "http://%s:%d/", ' % (_host, _port),
							'"format": "text"}',
						],
					),
				]
				reasoning = [
					"Two independent tool calls, ",
					"batched together.",
				]
			else:  # tool_call_unknown
				calls = [
					(
						"call_mock_3",
						"nonexistent_tool",
						['{"target": "test"}'],
					)
				]
				reasoning = ["Calling a tool the client ", "does not have."]
			self.log_message(
				"tool mode=%s calls=%s", mode, [c[1] for c in calls]
			)
			return self._stream(
				200, lambda w: _stream_tool_calls(w, cid, reasoning, calls)
			)

		if not stream:
			# Non-streaming success: single chat.completion object.
			content = " ".join(words)
			return self._send_json(
				200,
				{
					"id": cid,
					"object": "chat.completion",
					"created": _now(),
					"model": MODEL_ID,
					"choices": [
						{
							"index": 0,
							"message": {
								"role": "assistant",
								"content": content,
							},
							"finish_reason": "stop",
						}
					],
					"usage": {
						"prompt_tokens": 12,
						"completion_tokens": len(words),
						"total_tokens": 12 + len(words),
					},
				},
			)

		if mode == "error_during":
			return self._stream(
				200, lambda w: _stream_error_during(w, cid, words)
			)
		if mode == "no_done":
			return self._stream(200, lambda w: _stream_no_done(w, cid, words))
		return self._stream(200, lambda w: _stream_normal(w, cid, words))

	def log_message(self, format, *args):
		# Compact single-line logging to stdout.
		print(
			"%s - %s" % (self.log_date_time_string(), format % args), flush=True
		)


def main():
	ap = argparse.ArgumentParser(
		description="Mock llama-server for exercising zshell LLM error handling"
	)
	ap.add_argument("--port", type=int, default=8080)
	ap.add_argument("--host", default="127.0.0.1")
	args = ap.parse_args()

	global _host, _port
	_host = args.host
	_port = args.port

	server = ThreadingHTTPServer((args.host, args.port), Handler)
	server.daemon_threads = True
	print(
		f"mock llama-server on http://{args.host}:{args.port} "
		f"(model={MODEL_ID}, ctx={CTX_SIZE})\n"
		"type trigger words in the chat UI, or curl /_mock to change state",
		flush=True,
	)
	try:
		server.serve_forever()
	except KeyboardInterrupt:
		pass


if __name__ == "__main__":
	main()