"""
NEOME Bridge Client

Run examples:

	python bridgeclient.py --url http://127.0.0.1:11434 --backend ollama
	python bridgeclient.py --url http://127.0.0.1:1234 --backend lmstudio
	python bridgeclient.py --url http://127.0.0.1:8000 --backend vllm
	python bridgeclient.py --url http://127.0.0.1:8080 --backend localai
	python bridgeclient.py --url http://127.0.0.1:5000 --backend tgwui
	python bridgeclient.py --url http://127.0.0.1:5001 --backend koboldcpp
	python bridgeclient.py --url http://127.0.0.1:7860 --backend gradio
	python bridgeclient.py --url http://127.0.0.1:7861 --backend gradio
	python bridgeclient.py --url http://127.0.0.1:8188 --backend comfyui
	python bridgeclient.py --url http://127.0.0.1:3000 --backend custom
	python bridgeclient.py --url http://127.0.0.1:6006 --backend custom

If --backend is omitted, the script tries to auto-detect the backend.

Public API endpoint used by apps:

	https://bridge.neome.com/v1

Use the API key printed by this script:

	Authorization: Bearer sk_neome_xxxxx

OpenAI-style public routes:

	GET  /v1/models
	POST /v1/chat/completions
	POST /v1/completions
	POST /v1/embeddings
	POST /v1/images/generations
	POST /v1/images/edits
	POST /v1/images/variations
	POST /v1/audio/transcriptions
	POST /v1/audio/translations
	POST /v1/audio/speech
	POST /v1/videos/generations
	POST /v1/videos/edits

Custom backend:

	python bridgeclient.py --url http://127.0.0.1:3000 --backend custom

For custom backends, the bridge sends the original JSON payload and also adds:

	path = requested OpenAI-style path
	type = action name, such as chat, image, tts, stt, video

Example custom image request received by your local server:

	{
		"prompt": "a cute monkey",
		"path": "/v1/images/generations",
		"type": "image"
	}

The local backend only needs to return JSON. The bridge forwards it back to the public API response.
"""

import argparse
import asyncio
import json
import os
import secrets
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

import websockets

BRIDGE_WS = "wss://bridge.neome.com/ws"
PUBLIC_ENDPOINT = "https://bridge.neome.com/v1"
CONFIG_FILE = "bridge.json"

HTTP_TIMEOUT = 600
DETECT_TIMEOUT = 3
MAX_BODY_BYTES = 500 * 1024 * 1024

DEFAULT_URLS = [
	"http://127.0.0.1:11434",
	"http://127.0.0.1:1234",
	"http://127.0.0.1:8000",
	"http://127.0.0.1:8080",
	"http://127.0.0.1:5000",
	"http://127.0.0.1:5001",
	"http://127.0.0.1:7860",
	"http://127.0.0.1:7861",
	"http://127.0.0.1:8188",
	"http://127.0.0.1:3000",
	"http://127.0.0.1:6006",
]

BACKENDS = [
	"openai",
	"ollama",
	"a1111",
	"comfyui",
	"gradio",
	"omni",
	"piper",
	"kokoro",
	"coqui",
	"whisper",
	"localai",
	"lmstudio",
	"vllm",
	"llamacpp",
	"tgwui",
	"koboldcpp",
	"invokeai",
	"forge",
	"sdnext",
	"swarmui",
	"flux",
	"custom",
]


ACTIONS = {
	"/v1/models": "models",
	"/v1/chat/completions": "chat",
	"/v1/completions": "completion",
	"/v1/embeddings": "embedding",
	"/v1/images/generations": "image",
	"/v1/images/edits": "image_edit",
	"/v1/images/variations": "image_variation",
	"/v1/audio/transcriptions": "stt",
	"/v1/audio/translations": "translation",
	"/v1/audio/speech": "tts",
	"/v1/videos/generations": "video",
	"/v1/videos/edits": "video_edit",
}

ACTION_PATHS = {v: k for k, v in ACTIONS.items()}

ROUTES = {
	"/v1/models": {
		"*": "/v1/models",
		"ollama": "/api/tags",
		"a1111": "/sdapi/v1/options",
		"forge": "/sdapi/v1/options",
		"sdnext": "/sdapi/v1/options",
		"comfyui": "/system_stats",
		"gradio": "/config",
		"omni": "/",
		"piper": "/",
		"whisper": "/models",
		"custom": "/",
	},
	"/v1/chat/completions": {
		"*": "/v1/chat/completions",
		"ollama": "/api/chat",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/completions": {
		"*": "/v1/completions",
		"ollama": "/api/generate",
		"gradio": "/api/predict",
		"custom": "/",
	},
	"/v1/embeddings": {
		"*": "/v1/embeddings",
		"ollama": "/api/embeddings",
	},
	"/v1/images/generations": {
		"openai": "/v1/images/generations",
		"localai": "/v1/images/generations",
		"lmstudio": "/v1/images/generations",
		"vllm": "/v1/images/generations",
		"llamacpp": "/v1/images/generations",
		"flux": "/v1/images/generations",
		"a1111": "/sdapi/v1/txt2img",
		"forge": "/sdapi/v1/txt2img",
		"sdnext": "/sdapi/v1/txt2img",
		"invokeai": "/api/v1/images",
		"swarmui": "/GenerateText2Image",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/images/edits": {
		"openai": "/v1/images/edits",
		"localai": "/v1/images/edits",
		"flux": "/v1/images/edits",
		"a1111": "/sdapi/v1/img2img",
		"forge": "/sdapi/v1/img2img",
		"sdnext": "/sdapi/v1/img2img",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/images/variations": {
		"openai": "/v1/images/variations",
		"localai": "/v1/images/variations",
		"a1111": "/sdapi/v1/img2img",
		"forge": "/sdapi/v1/img2img",
		"sdnext": "/sdapi/v1/img2img",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/audio/transcriptions": {
		"openai": "/v1/audio/transcriptions",
		"localai": "/v1/audio/transcriptions",
		"whisper": "/inference",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/audio/translations": {
		"openai": "/v1/audio/translations",
		"localai": "/v1/audio/translations",
		"whisper": "/inference",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/audio/speech": {
		"openai": "/v1/audio/speech",
		"localai": "/v1/audio/speech",
		"kokoro": "/v1/audio/speech",
		"omni": "/tts_to_audio/",
		"piper": "/tts",
		"coqui": "/api/tts",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/videos/generations": {
		"openai": "/v1/videos/generations",
		"localai": "/v1/videos/generations",
		"wan": "/videos/generations",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
	"/v1/videos/edits": {
		"openai": "/v1/videos/edits",
		"wan": "/videos/edits",
		"gradio": "/api/predict",
		"comfyui": "/prompt",
		"custom": "/",
	},
}

OPENAI_BACKENDS = {
	"openai",
	"localai",
	"lmstudio",
	"vllm",
	"llamacpp",
	"tgwui",
	"koboldcpp",
	"kokoro",
	"invokeai",
	"swarmui",
}

A1111_BACKENDS = {
	"a1111",
	"forge",
	"sdnext",
}

running = True


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


def make_key():
	return "sk_neome_" + secrets.token_urlsafe(32)


def load_config():
	if not os.path.exists(CONFIG_FILE):
		return {}

	try:
		with open(CONFIG_FILE, "r", encoding="utf-8") as f:
			data = json.load(f)

		return data if isinstance(data, dict) else {}

	except Exception:
		return {}


def save_config(config):
	tmp = CONFIG_FILE + ".tmp"

	with open(tmp, "w", encoding="utf-8") as f:
		json.dump(config, f, indent=2)

	os.replace(tmp, CONFIG_FILE)


def safe_url(url):
	p = urlparse(url)

	if p.scheme not in ("http", "https"):
		raise ValueError("URL must start with http:// or https://")

	if not p.hostname:
		raise ValueError("URL is missing hostname")

	return url.rstrip("/")


def build_url(base_url, path):
	base_url = safe_url(base_url)
	p = urlparse(base_url)

	origin = f"{p.scheme}://{p.netloc}"
	base_path = p.path.rstrip("/")
	path = "/" + str(path or "").lstrip("/")

	if base_path and path.startswith(base_path + "/"):
		return origin + path

	return base_url + path


def http_json(method, base_url, path, payload=None, timeout=HTTP_TIMEOUT):
	base_url = safe_url(base_url)
	data = None

	headers = {
		"Content-Type": "application/json",
		"Accept": "application/json",
	}

	if payload is not None:
		data = json.dumps(payload).encode("utf-8")

	req = Request(
		build_url(base_url, path),
		data=data,
		headers=headers,
		method=method.upper(),
	)

	try:
		with urlopen(req, timeout=timeout) as r:
			raw = r.read(MAX_BODY_BYTES + 1)

			if len(raw) > MAX_BODY_BYTES:
				raise RuntimeError("Local API response too large")

			if not raw:
				return {}

			text = raw.decode("utf-8", errors="replace")

			try:
				return json.loads(text)
			except Exception:
				return {
					"raw": text,
				}

	except HTTPError as e:
		raw = e.read(MAX_BODY_BYTES).decode("utf-8", errors="replace")

		try:
			body = json.loads(raw)
		except Exception:
			body = {
				"error": {
					"message": raw or str(e),
					"type": "local_http_error",
					"status": e.code,
				}
			}

		raise RuntimeError(json.dumps(body))

	except URLError as e:
		raise RuntimeError(str(e.reason))


def http_get(url, path, timeout=DETECT_TIMEOUT):
	return http_json("GET", url, path, None, timeout)


def http_post(url, path, payload, timeout=HTTP_TIMEOUT):
	return http_json("POST", url, path, payload, timeout)


def detect_ollama(url):
	data = http_get(url, "/api/tags")
	models = [m.get("name") for m in data.get("models", []) if m.get("name")]

	return {
		"backend": "ollama",
		"url": url,
		"models": models,
		"default_model": models[0] if models else "default",
	}


def detect_a1111(url, backend="a1111"):
	http_get(url, "/sdapi/v1/options")

	return {
		"backend": backend,
		"url": url,
		"models": ["default"],
		"default_model": "default",
	}


def detect_comfyui(url):
	http_get(url, "/system_stats")

	return {
		"backend": "comfyui",
		"url": url,
		"models": ["default"],
		"default_model": "default",
	}


def detect_gradio(url):
	data = http_get(url, "/config")

	if not isinstance(data, dict):
		raise RuntimeError("Not Gradio")

	if "dependencies" not in data and "components" not in data:
		raise RuntimeError("Not Gradio")

	return {
		"backend": "gradio",
		"url": url,
		"models": ["default"],
		"default_model": "default",
		"config": data,
	}


def detect_omni(url):
	data = http_get(url, "/")

	text = json.dumps(data).lower()

	if "omnivoice" not in text and "tts_to_audio" not in text:
		raise RuntimeError("Not OmniVoice")

	return {
		"backend": "omni",
		"url": url,
		"models": ["omnivoice"],
		"default_model": "omnivoice",
	}

def detect_piper(url):
	data = http_get(url, "/")

	text = json.dumps(data).lower()

	if "piper" not in text and "tts" not in text:
		raise RuntimeError("Not Piper")

	return {
		"backend": "piper",
		"url": url,
		"models": ["piper"],
		"default_model": "piper",
	}


def detect_whisper(url):
	try:
		data = http_get(url, "/models")
	except Exception:
		data = http_get(url, "/")

	text = json.dumps(data).lower()

	if "whisper" not in text and "transcription" not in text:
		raise RuntimeError("Not Whisper")

	return {
		"backend": "whisper",
		"url": url,
		"models": ["whisper"],
		"default_model": "whisper",
	}


def detect_openai(url, backend="openai"):
	data = http_get(url, "/v1/models")
	models = [m.get("id") for m in data.get("data", []) if m.get("id")]

	return {
		"backend": backend,
		"url": url,
		"models": models,
		"default_model": models[0] if models else "default",
	}


def detect_custom(url):
	try:
		http_get(url, "/")
	except Exception:
		http_get(url, "/health")

	return {
		"backend": "custom",
		"url": url,
		"models": ["default"],
		"default_model": "default",
	}


def guess_backend_from_url(url):
	p = urlparse(url)
	port = str(p.port or "")
	host = str(p.hostname or "").lower()

	if "omni" in host:
		return "omni"

	if "image.neome.com" in host:
		return "openai"

	if port == "11434":
		return "ollama"

	if port == "1234":
		return "lmstudio"

	if port == "8000":
		return "vllm"

	if port == "8080":
		return "localai"

	if port == "5000":
		return "tgwui"

	if port == "5001":
		return "koboldcpp"

	if port in ("7860", "7861"):
		return "gradio"

	if port == "8188":
		return "comfyui"

	return "openai"


def detect_forced(url, forced_backend):
	if forced_backend == "omni":
		return {
			"backend": "omni",
			"url": url,
			"models": ["omnivoice"],
			"default_model": "omnivoice",
		}

	if forced_backend == "ollama":
		return detect_ollama(url)

	if forced_backend == "a1111":
		return detect_a1111(url, "a1111")

	if forced_backend == "forge":
		return detect_a1111(url, "forge")

	if forced_backend == "sdnext":
		return detect_a1111(url, "sdnext")

	if forced_backend == "comfyui":
		return detect_comfyui(url)

	if forced_backend == "gradio":
		return detect_gradio(url)

	if forced_backend == "piper":
		return {
			"backend": "piper",
			"url": url,
			"models": ["piper"],
			"default_model": "piper",
		}

	if forced_backend == "whisper":
		return {
			"backend": "whisper",
			"url": url,
			"models": ["whisper"],
			"default_model": "whisper",
		}

	if forced_backend == "custom":
		return {
			"backend": "custom",
			"url": url,
			"models": ["default"],
			"default_model": "default",
		}

	if forced_backend in BACKENDS:
		return detect_openai(url, forced_backend)

	raise RuntimeError("Unknown backend: " + forced_backend)


def detect_backend(url, forced_backend=""):
	if forced_backend:
		return detect_forced(url, forced_backend)

	guessed = guess_backend_from_url(url)

	if guessed == "omni":
		return {
			"backend": "omni",
			"url": url,
			"models": ["omnivoice"],
			"default_model": "omnivoice",
		}

	if guessed == "ollama":
		try:
			return detect_ollama(url)
		except Exception:
			pass

	if guessed == "gradio":
		try:
			return detect_gradio(url)
		except Exception:
			pass

	if guessed == "comfyui":
		try:
			return detect_comfyui(url)
		except Exception:
			pass

	for fn in (
		detect_omni,
		detect_ollama,
		detect_a1111,
		detect_comfyui,
		detect_gradio,
	):
		try:
			return fn(url)
		except Exception:
			pass

	try:
		return detect_openai(url, guessed)
	except Exception:
		pass

	return None


def auto_detect(custom_url="", forced_backend=""):
	urls = []

	if custom_url:
		urls.append(custom_url.rstrip("/"))

	for url in DEFAULT_URLS:
		url = url.rstrip("/")
		if url not in urls:
			urls.append(url)

	for url in urls:
		if not url:
			continue

		try:
			info = detect_backend(url, forced_backend)

			if info:
				return info

		except Exception:
			pass

	return None


def pick_model(payload, info):
	requested = str(payload.get("model") or "") if isinstance(payload, dict) else ""
	models = info.get("models") or []
	default_model = info.get("default_model") or ""

	if requested and requested != "auto" and requested in models:
		return requested

	if default_model:
		return default_model

	if models:
		return models[0]

	return requested or "default"


def strip_data_url(value):
	value = str(value or "")

	if value.startswith("data:") and "," in value:
		return value.split(",", 1)[1]

	return value


def get_images(payload):
	images = payload.get("image") or payload.get("images") or []

	if isinstance(images, str):
		images = [images]

	return [strip_data_url(x) for x in images if x]


def parse_size(payload):
	size = str(payload.get("size") or "1024x1024").lower()
	width = int(payload.get("width") or 1024)
	height = int(payload.get("height") or 1024)

	if "x" in size:
		try:
			w, h = size.split("x", 1)
			width = int(w)
			height = int(h)
		except Exception:
			pass

	return width, height


def route_for(path, backend):
	items = ROUTES.get(path) or {}

	return items.get(backend) or items.get("*") or ""


def normalize_path(path, body):
	if path in ("/v1", "/v1/"):
		action = str(body.get("type") or "chat")
		return ACTION_PATHS.get(action, "/v1/chat/completions")

	return path


def model_list(info):
	models = info.get("models") or ["default"]

	return {
		"object": "list",
		"data": [
			{
				"id": name,
				"object": "model",
				"created": now(),
				"owned_by": info.get("backend") or "local",
			}
			for name in models
		],
	}


def openai_payload(payload, info):
	if isinstance(payload, dict) and "model" in payload:
		payload = dict(payload)
		payload["model"] = pick_model(payload, info)

	return payload or {}


def ollama_messages(payload):
	messages = payload.get("messages")

	if isinstance(messages, list):
		out = []

		for msg in messages:
			item = dict(msg)
			content = item.get("content")

			if isinstance(content, list):
				texts = []
				images = []

				for part in content:
					if part.get("type") == "text":
						texts.append(str(part.get("text") or ""))

					if part.get("type") == "image_url":
						url = ((part.get("image_url") or {}).get("url") or "")

						if url:
							images.append(strip_data_url(url))

				item["content"] = "\n".join(texts)

				if images:
					item["images"] = images

			out.append(item)

		return out

	prompt = payload.get("prompt") or payload.get("input") or ""

	return [
		{
			"role": "user",
			"content": str(prompt),
		}
	]


def ollama_payload(path, payload, info):
	model = pick_model(payload, info)

	if path == "/v1/chat/completions":
		out = {
			"model": model,
			"messages": ollama_messages(payload),
			"stream": False,
		}

		images = get_images(payload)

		if images:
			for msg in reversed(out["messages"]):
				if msg.get("role") == "user":
					msg["images"] = images
					break

		options = {}

		for key in ("temperature", "top_p", "top_k", "seed", "num_ctx"):
			if key in payload:
				options[key] = payload[key]

		if "max_tokens" in payload:
			options["num_predict"] = payload["max_tokens"]

		if options:
			out["options"] = options

		return out

	if path == "/v1/completions":
		prompt = payload.get("prompt") or payload.get("input") or ""

		if isinstance(prompt, list):
			prompt = "\n".join(str(x) for x in prompt)

		out = {
			"model": model,
			"prompt": str(prompt),
			"stream": False,
		}

		images = get_images(payload)

		if images:
			out["images"] = images

		return out

	if path == "/v1/embeddings":
		return {
			"model": model,
			"prompt": str(payload.get("input") or ""),
		}

	return payload


def a1111_payload(path, payload, info):
	width, height = parse_size(payload)

	if path == "/v1/images/generations":
		return {
			"prompt": payload.get("prompt", ""),
			"negative_prompt": payload.get("negative_prompt", ""),
			"steps": payload.get("steps", 20),
			"width": width,
			"height": height,
			"cfg_scale": payload.get("cfg_scale", 7),
			"seed": payload.get("seed", -1),
			"batch_size": payload.get("n", 1),
		}

	if path in ("/v1/images/edits", "/v1/images/variations"):
		images = get_images(payload)

		return {
			"prompt": payload.get("prompt", ""),
			"negative_prompt": payload.get("negative_prompt", ""),
			"init_images": images,
			"denoising_strength": payload.get("strength", 0.75),
			"steps": payload.get("steps", 20),
			"width": width,
			"height": height,
			"cfg_scale": payload.get("cfg_scale", 7),
			"seed": payload.get("seed", -1),
			"batch_size": payload.get("n", 1),
		}

	return payload


def omni_payload(path, payload, info):
	if path == "/v1/audio/speech":
		return {
			"text": payload.get("input") or payload.get("text") or "",
			"voice": payload.get("voice") or "default",
			"format": payload.get("response_format") or payload.get("format") or "mp3",
		}

	return payload or {}


def piper_payload(path, payload, info):
	if path == "/v1/audio/speech":
		return {
			"text": payload.get("input") or payload.get("text") or "",
			"voice": payload.get("voice") or "default",
		}

	return payload or {}


def gradio_payload(path, payload, info):
	prompt = (
		payload.get("prompt") or
		payload.get("input") or
		payload.get("text") or
		""
	)

	if path == "/v1/chat/completions":
		messages = payload.get("messages") or []

		if messages:
			last = messages[-1]
			prompt = last.get("content") or prompt

			if isinstance(prompt, list):
				parts = []

				for part in prompt:
					if part.get("type") == "text":
						parts.append(str(part.get("text") or ""))

				prompt = "\n".join(parts)

	if path == "/v1/audio/speech":
		prompt = payload.get("input") or payload.get("text") or ""

	return {
		"data": [
			prompt,
		]
	}


def custom_payload(path, payload, info):
	out = dict(payload or {})
	out.setdefault("path", path)
	out.setdefault("type", ACTIONS.get(path, "custom"))

	return out


def convert_payload(path, payload, info):
	backend = info.get("backend")

	if backend == "ollama":
		return ollama_payload(path, payload, info)

	if backend in A1111_BACKENDS:
		return a1111_payload(path, payload, info)

	if backend == "omni":
		return omni_payload(path, payload, info)

	if backend == "piper":
		return piper_payload(path, payload, info)

	if backend == "gradio":
		return gradio_payload(path, payload, info)

	if backend == "custom":
		return custom_payload(path, payload, info)

	if backend in OPENAI_BACKENDS:
		return openai_payload(payload, info)

	return payload or {}


def normalize_gradio_text(data):
	if isinstance(data, dict):
		if "data" in data and isinstance(data["data"], list) and data["data"]:
			return str(data["data"][0])

		if "output" in data:
			return str(data["output"])

		if "text" in data:
			return str(data["text"])

	if isinstance(data, list) and data:
		return str(data[0])

	return str(data)


def convert_response(path, data, info):
	backend = info.get("backend")

	if backend == "ollama":
		if path == "/v1/models":
			models = [m.get("name") for m in data.get("models", []) if m.get("name")]
			info["models"] = models

			return model_list(info)

		if path == "/v1/chat/completions":
			content = ""

			if isinstance(data.get("message"), dict):
				content = data["message"].get("content", "")

			return {
				"id": "chatcmpl_" + secrets.token_hex(12),
				"object": "chat.completion",
				"created": now(),
				"model": info.get("default_model", "default"),
				"choices": [
					{
						"index": 0,
						"message": {
							"role": "assistant",
							"content": content,
						},
						"finish_reason": "stop",
					}
				],
			}

		if path == "/v1/completions":
			return {
				"id": "cmpl_" + secrets.token_hex(12),
				"object": "text_completion",
				"created": now(),
				"model": info.get("default_model", "default"),
				"choices": [
					{
						"text": data.get("response", ""),
						"index": 0,
						"logprobs": None,
						"finish_reason": "stop",
					}
				],
			}

		if path == "/v1/embeddings":
			return {
				"object": "list",
				"data": [
					{
						"object": "embedding",
						"index": 0,
						"embedding": data.get("embedding", []),
					}
				],
			}

	if backend in A1111_BACKENDS:
		if path == "/v1/models":
			return model_list(info)

		if path.startswith("/v1/images/"):
			return {
				"created": now(),
				"data": [
					{
						"b64_json": img,
					}
					for img in data.get("images", [])
				],
			}

	if backend == "gradio":
		text = normalize_gradio_text(data)

		if path == "/v1/chat/completions":
			return {
				"id": "chatcmpl_" + secrets.token_hex(12),
				"object": "chat.completion",
				"created": now(),
				"model": info.get("default_model", "default"),
				"choices": [
					{
						"index": 0,
						"message": {
							"role": "assistant",
							"content": text,
						},
						"finish_reason": "stop",
					}
				],
			}

		if path == "/v1/completions":
			return {
				"id": "cmpl_" + secrets.token_hex(12),
				"object": "text_completion",
				"created": now(),
				"model": info.get("default_model", "default"),
				"choices": [
					{
						"text": text,
						"index": 0,
						"logprobs": None,
						"finish_reason": "stop",
					}
				],
			}

	if backend in ("omni", "piper") and path == "/v1/audio/speech":
		return data

	if path == "/v1/models":
		return model_list(info)

	return data


def unsupported(path, backend):
	return {
		"error": {
			"message": "Route not supported by this backend.",
			"type": "unsupported_route",
			"path": path,
			"backend": backend,
		}
	}


def handle_request(info, req):
	method = str(req.get("method") or "POST").upper()
	body = req.get("body") or {}
	path = normalize_path(str(req.get("path") or ""), body)
	backend = info.get("backend")
	local_route = route_for(path, backend)

	if not local_route:
		return unsupported(path, backend)

	if path == "/v1/models" and backend in A1111_BACKENDS.union({"comfyui", "omni", "piper", "gradio", "custom"}):
		return model_list(info)

	payload = convert_payload(path, body, info)

	if path == "/v1/models":
		data = http_get(info["url"], local_route, HTTP_TIMEOUT)
	elif method == "GET":
		data = http_get(info["url"], local_route, HTTP_TIMEOUT)
	else:
		data = http_post(info["url"], local_route, payload, HTTP_TIMEOUT)

	return convert_response(path, data, info)


async def send_json(ws, data):
	await ws.send(json.dumps(data, separators=(",", ":")))


async def register(ws, key):
	await send_json(ws, {
		"type": "register",
		"key": key,
	})


async def process_message(ws, info, message):
	try:
		req = json.loads(message)
	except Exception:
		return

	if req.get("type") != "request":
		return

	request_id = req.get("request_id")

	try:
		body = await asyncio.to_thread(handle_request, info, req)
		status = 200

		if isinstance(body, dict) and "error" in body:
			status = 400

		await send_json(ws, {
			"type": "response",
			"request_id": request_id,
			"status": status,
			"body": body,
		})

	except Exception as e:
		await send_json(ws, {
			"type": "response",
			"request_id": request_id,
			"status": 500,
			"body": {
				"error": {
					"message": str(e),
					"type": "bridge_error",
				}
			},
		})


async def connect_loop(ws_url, key, info):
	global running
	backoff = 1

	while running:
		try:
			print("Connecting...", flush=True)

			async with websockets.connect(
				ws_url,
				ping_interval=25,
				ping_timeout=10,
				max_size=None,
				close_timeout=5,
			) as ws:
				await register(ws, key)
				backoff = 1

				print("Connected ✓")
				print("")
				print("API KEY:")
				print(key)
				print("")
				print("Endpoint:")
				print(PUBLIC_ENDPOINT)
				print("")
				print("Detected:")
				print(f"{info.get('backend')} ({info.get('url')})")
				print("")
				print("Default model:")
				print(info.get("default_model") or "default")
				print("")
				print("Ready.", flush=True)

				async for message in ws:
					await process_message(ws, info, message)

		except Exception as e:
			if running:
				print("Disconnected:", str(e), flush=True)
				print(f"Reconnecting in {backoff}s...", flush=True)
				await asyncio.sleep(backoff)
				backoff = min(backoff * 2, 30)


async def main():
	parser = argparse.ArgumentParser()
	parser.add_argument("--url", default="", help="Local API URL")
	parser.add_argument("--key", default="", help="Existing bridge key")
	parser.add_argument("--ws", default=BRIDGE_WS, help="Bridge WebSocket URL")
	parser.add_argument("--backend", default="", help="Force backend name")
	args = parser.parse_args()

	config = load_config()
	loaded_config = bool(config.get("key"))
	key = (
		args.key
		or config.get("key")
		or make_key()
	)
	local_url = args.url or ""
	forced_backend = args.backend or ""

	try:
		if local_url:
			local_url = safe_url(local_url)
	except Exception as e:
		print("Invalid URL:", e, flush=True)
		sys.exit(1)

	info = auto_detect(local_url, forced_backend)

	if not info:
		print("No local AI backend found.")
		print("Try one of:")

		for url in DEFAULT_URLS:
			print(url)

		print("")
		print("Or run:")
		print("python bridgeclient.py --url http://127.0.0.1:11434 --backend ollama")
		print("python bridgeclient.py --url https://omni.neome.com --backend omni")
		print("python bridgeclient.py --url http://127.0.0.1:7860 --backend gradio")
		sys.exit(1)

	config["key"] = key
	config["url"] = info["url"]
	config["backend"] = info["backend"]
	save_config(config)

	print("=========================")
	print("NEOME BRIDGE")
	print("=========================")
	if loaded_config:
		print("Loaded config: bridge.json")
	else:
		print("Created new bridge.json")
	print("")
	print("Config file:")
	print(os.path.abspath(CONFIG_FILE))
	print("")
	await connect_loop(args.ws, key, info)


if __name__ == "__main__":
	try:
		asyncio.run(main())
	except KeyboardInterrupt:
		running = False