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

MODEL_ID = "mock-model"
CTX_SIZE = 8192

_host = "127.0.0.1"
_port = 8080

_state_lock = threading.Lock()
_state = {
	"loading": False,
	"no_models": False,
	"chat_mode": None,
}


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):
	return ("data: " + json.dumps(obj, separators=(",", ":")) + "\n\n").encode()


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


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"
	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):
	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):
	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)
	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)


def _last_user_text(messages):
	for msg in reversed(messages or []):
		if msg.get("role") == "user":
			content = msg.get("content")
			if isinstance(content, list):
				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"

	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):
		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:
			with contextlib.suppress(BrokenPipeError, ConnectionResetError):
				self.wfile.flush()

	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"))

	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})

	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":
			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":
			return self._send_json(
				200, _error(400, "Mock: 200-with-json-error", "server_error")
			)

		if mode in _TOOL_MODES:
			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",
						[
							f'{{"url": "http://{_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",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							f'{{"url": "http://{_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",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_1",
						"readfile",
						['{"path": "/etc/', 'hostname"}'],
					),
					(
						"call_mock_2",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_3",
						"readfile",
						['{"path": "/etc/', 'os-release"}'],
					),
					(
						"call_mock_4",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
					(
						"call_mock_5",
						"readfile",
						['{"path": "/proc/', 'uptime"}'],
					),
					(
						"call_mock_6",
						"webfetch",
						[
							f'{{"url": "http://{_host}:{_port}/", ',
							'"format": "text"}',
						],
					),
				]
				reasoning = [
					"Two independent tool calls, ",
					"batched together.",
				]
			else:
				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:
			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):
		print(f"{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,
	)
	with contextlib.suppress(KeyboardInterrupt):
		server.serve_forever()


if __name__ == "__main__":
	main()