import asyncio
import websockets
import os

# --- Configuration ---
HTTP_PORT = 7860
WEBSOCKET_PORT = 8765
MINECRAFT_RAM = "12G"
PLAYIT_SECRET = os.environ.get('PLAYIT_SECRET')

# A set to keep track of all connected WebSocket clients
connected_clients = set()
# Global reference to the Minecraft process
mc_process = None

async def start_subprocess(command):
    """Starts a subprocess and returns the process object."""
    print(f"Starting subprocess: {' '.join(command)}")
    process = await asyncio.create_subprocess_exec(
        *command,
        stdout=asyncio.subprocess.PIPE,
        stdin=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    return process

async def broadcast(message):
    """Sends a message to all connected clients."""
    if connected_clients:
        # Use asyncio.create_task to send messages concurrently
        tasks = [asyncio.create_task(client.send(message)) for client in connected_clients]
        # Allow other tasks to run while sending
        await asyncio.sleep(0)

async def read_stream(stream, prefix=""):
    """Reads lines from a stream and broadcasts them."""
    while True:
        line = await stream.readline()
        if not line:
            await broadcast(f"[SYSTEM] Stream '{prefix}' closed.")
            break
        output = line.decode('utf-8', errors='replace').strip()
        print(f"{prefix}{output}")
        await broadcast(f"{prefix}{output}")

async def run_minecraft():
    """The main task for running and monitoring the Minecraft server."""
    global mc_process
    await broadcast("[SYSTEM] Starting Minecraft server process... (This may take a few minutes on first run)")
    
    mc_command = ['java', f'-Xmx{MINECRAFT_RAM}', '-Xms1G', '-jar', 'server.jar', 'nogui']
    mc_process = await start_subprocess(mc_command)
    
    # Create tasks to forward Minecraft's output to websockets
    stdout_task = asyncio.create_task(read_stream(mc_process.stdout, ""))
    stderr_task = asyncio.create_task(read_stream(mc_process.stderr, "[mc/err] "))

    # Wait for the process to finish
    exit_code = await mc_process.wait()
    print(f"Minecraft process has exited with code {exit_code}.")
    await broadcast(f"[SYSTEM] Minecraft server has stopped with exit code {exit_code}.")
    await asyncio.gather(stdout_task, stderr_task)


async def websocket_handler(websocket, path):
    """Handles WebSocket connections."""
    global mc_process
    print("Client connected to WebSocket.")
    connected_clients.add(websocket)
    try:
        async for message in websocket:
            print(f"Received command: {message}")
            if mc_process and mc_process.returncode is None:
                command = message.strip() + '\n'
                mc_process.stdin.write(command.encode('utf-8'))
                await mc_process.stdin.drain()
            else:
                await websocket.send("Error: Minecraft server is not running or has stopped.")
    except websockets.exceptions.ConnectionClosed:
        print("Client disconnected.")
    finally:
        connected_clients.remove(websocket)

async def http_handler(reader, writer):
    """Handles basic HTTP requests to serve the HTML file."""
    try:
        request_line_bytes = await reader.readline()
        if not request_line_bytes or request_line_bytes.isspace():
            return
        
        request_line = request_line_bytes.decode()
        method, path, _ = request_line.split(' ', 2)
        
        if method == "GET" and path == "/":
            writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n")
            with open("index.html", "rb") as f:
                writer.write(f.read())
        else:
            writer.write(b"HTTP/1.1 404 Not Found\r\n\r\nNot Found")
            
        await writer.drain()
    finally:
        if not writer.is_closing():
            writer.close()
            await writer.wait_closed()

async def main():
    # --- THIS IS THE NEW STRUCTURE ---

    # 1. Start the web servers IMMEDIATELY. This is the key.
    http_server = await asyncio.start_server(http_handler, "0.0.0.0", HTTP_PORT)
    websocket_server = await websockets.serve(websocket_handler, "0.0.0.0", WEBSOCKET_PORT)
    print(f"HTTP server started on port {HTTP_PORT}. APPLICATION IS INITIALIZED.")
    print(f"WebSocket server started on port {WEBSOCKET_PORT}.")

    # 2. Start playit.gg in the background.
    if PLAYIT_SECRET:
        playit_process = await start_subprocess(['playit', '--secret', PLAYIT_SECRET])
        asyncio.create_task(read_stream(playit_process.stdout, "[playit] "))
        asyncio.create_task(read_stream(playit_process.stderr, "[playit/err] "))
    else:
        print("Warning: PLAYIT_SECRET not set.")

    # 3. Start the slow Minecraft server in the background.
    asyncio.create_task(run_minecraft())

    # 4. Wait forever. This keeps the script, and thus the servers and container, alive.
    await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("Server shutting down.")