import asyncio
import json
import os
import sqlite3
import websockets


# =========================
# CONFIG
# =========================

SERVER_URL = "wss://env.neome.com"

DB_PATH = "./database.db"
DB_PASSWORD = "CHANGE_ME"

WS_MAX_SIZE = 50 * 1024 * 1024


# =========================
# DB
# =========================

def get_conn():
	return sqlite3.connect(DB_PATH)


def run_query(sql):
	conn = get_conn()

	try:
		cursor = conn.cursor()
		cursor.execute(sql)

		rows = cursor.fetchall()

		if not rows:
			return "[empty]"

		out = []

		for row in rows:
			out.append(" | ".join(str(x) for x in row))

		return "\n".join(out)

	finally:
		conn.close()


def run_write(sql):
	conn = get_conn()

	try:
		cursor = conn.cursor()
		cursor.execute(sql)
		conn.commit()

		return "OK"

	finally:
		conn.close()


def get_schema():
	conn = get_conn()

	try:
		cursor = conn.cursor()

		cursor.execute("""
			SELECT name
			FROM sqlite_master
			WHERE type='table'
			ORDER BY name
		""")

		tables = cursor.fetchall()

		if not tables:
			return "[empty]"

		out = []

		for table_row in tables:
			table = table_row[0]

			out.append("")
			out.append("TABLE: " + table)

			cursor.execute("PRAGMA table_info(" + table + ")")

			columns = cursor.fetchall()

			for col in columns:
				out.append("  " + str(col[1]) + " (" + str(col[2]) + ")")

		return "\n".join(out).strip()

	finally:
		conn.close()


# =========================
# COMMANDS
# =========================

def handle_command(cmd):
	try:
		if not isinstance(cmd, list) or not cmd:
			return "Invalid command"

		cmd_type = str(cmd[0] or "").lower()

		if cmd_type == "schema":
			return get_schema()

		elif cmd_type == "query":
			if len(cmd) < 2:
				return "Missing SQL"

			return run_query(str(cmd[1]))

		elif cmd_type == "write":
			if len(cmd) < 2:
				return "Missing SQL"

			return run_write(str(cmd[1]))

		return "Unknown command: " + cmd_type

	except Exception as e:
		return "Error: " + str(e)


# =========================
# SOCKET
# =========================

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


async def main():
	while True:
		try:
			async with websockets.connect(
				SERVER_URL,
				max_size=WS_MAX_SIZE
			) as ws:

				await send_json(ws, {
					"role": "client"
				})

				print("[+] connected", flush=True)

				async for raw in ws:
					try:
						data = json.loads(raw)
					except:
						continue

					if data.get("type") == "hello":
						print("CLIENT_ID:", data.get("id"), flush=True)
						continue

					if data.get("type") == "cmd":
						result = handle_command(
							data.get("cmd")
						)

						await send_json(ws, {
							"type": "result",
							"output": result
						})

		except Exception as e:
			print("[!] disconnected:", e, flush=True)
			await asyncio.sleep(2)


if __name__ == "__main__":
	asyncio.run(main())