import asyncio
import websockets
import json
import random
import string
import agentv2
import http.server
import threading
import socketserver

# =========================================================
# GLOBAL ROOM STATE (single room game)
# =========================================================

ROOM = {
    "code": None,
    "secret": None,
    "difficulty": 1,
    "judge": "lenient",
    "tries": 10,
    "is_active": False,
    "admin_ws": None,
    "players": {}
}


# =========================================================
# UTIL HELPERS
# =========================================================

async def send(ws, payload: dict):
    """Safe JSON send."""
    try:
        await ws.send(json.dumps(payload))
    except:
        pass


async def broadcast(payload: dict, exclude=None):
    """Broadcast to all players safely."""
    dead = []

    for ws in ROOM["players"]:
        if ws == exclude:
            continue
        try:
            await ws.send(json.dumps(payload))
        except:
            dead.append(ws)

    # cleanup dead sockets
    for ws in dead:
        ROOM["players"].pop(ws, None)


def reset_room():
    """Fully reset room state."""
    ROOM.update({
        "code": None,
        "secret": None,
        "difficulty": 1,
        "judge": "lenient",
        "tries": 10,
        "is_active": False,
        "players": {}
    })


# =========================================================
# ADMIN HANDLER
# =========================================================

async def handle_admin(websocket):
    print("👑 Admin Connected")
    ROOM["admin_ws"] = websocket

    try:
        async for message in websocket:
            data = json.loads(message)

            # -----------------------------
            # START SESSION
            # -----------------------------
            if data["type"] == "admin_start":

                ROOM["code"] = ''.join(
                    random.choices(string.ascii_uppercase + string.digits, k=4)
                )

                ROOM["secret"] = data["secret"]
                ROOM["difficulty"] = data["diff"]
                ROOM["judge"] = data["judge"]
                ROOM["tries"] = int(data["tries"])
                ROOM["is_active"] = True

                print(f"🚀 Room Created: {ROOM['code']}")

                await send(websocket, {
                    "type": "room_created",
                    "code": ROOM["code"],
                    "secret": ROOM["secret"]
                })

            # -----------------------------
            # END SESSION
            # -----------------------------
            elif data["type"] == "admin_end":

                ROOM["is_active"] = False

                await broadcast({
                    "type": "system_notice",
                    "message": "SESSION TERMINATED BY ADMIN."
                })

                reset_room()

    finally:
        ROOM["admin_ws"] = None


# =========================================================
# PLAYER LOGIC HELPERS
# =========================================================

async def send_admin_chat(session, player, message):
    if ROOM["admin_ws"]:
        await send(ROOM["admin_ws"], {
            "type": "live_chat",
            "session": session,
            "player": player,
            "message": message
        })


async def handle_hint(websocket, player_name, player_data):
    if player_data["attempts"] + 2 > ROOM["tries"]:
        await send(websocket, {
            "type": "game_over",
            "result": "loss",
            "message": "INSUFFICIENT ATTEMPTS FOR HINT. SYSTEM LOCKOUT."
        })
        return False

    player_data["attempts"] += 2

    hint = agentv2.generate_hint(ROOM["secret"])
    msg = f"[TACTICAL HINT]: {hint} (Cost: 2 Attempts)"

    await send(websocket, {
        "type": "chat_response",
        "sender": "System",
        "message": msg
    })

    await send_admin_chat(player_name, "System", f"PURCHASED HINT: {hint}")

    return True


async def handle_ai_response(websocket, history, player_name):
    ai_reply = agentv2.generate_response(
        history,
        ROOM["secret"],
        ROOM["difficulty"],
        ROOM["judge"]
    )

    history.append({"role": "assistant", "content": ai_reply})

    await send(websocket, {
        "type": "chat_response",
        "sender": "Agent",
        "message": ai_reply
    })

    await send_admin_chat(player_name, "AI", ai_reply)


# =========================================================
# PLAYER HANDLER
# =========================================================

async def handle_player(websocket):

    player_name = "Unknown"

    try:
        join_msg = json.loads(await websocket.recv())

        if join_msg.get("type") != "join":
            return

        if not ROOM["is_active"]:
            await send(websocket, {"type": "error", "message": "SESSION CLOSED."})
            return

        if join_msg["code"].upper() != ROOM["code"]:
            await send(websocket, {"type": "error", "message": "INVALID ROOM CODE"})
            return

        # -----------------------------
        # JOIN
        # -----------------------------
        player_name = join_msg["name"]

        ROOM["players"][websocket] = {
            "name": player_name,
            "attempts": 0
        }

        print(f"👤 Player Joined: {player_name}")

        await send(websocket, {
            "type": "joined",
            "diff": ROOM["difficulty"],
            "tries": ROOM["tries"]
        })

        if ROOM["admin_ws"]:
            # Send updated list to Admin
            names = [p["name"] for p in ROOM["players"].values()]
            await send(ROOM["admin_ws"], {
                "type": "player_update",
                "count": len(ROOM["players"]),
                "names": names
            })

        await broadcast({
            "type": "system_notice",
            "message": f"OPERATIVE '{player_name}' CONNECTED."
        })

        history = []
        secret_upper = ROOM["secret"].upper()

        # -----------------------------
        # CHAT LOOP
        # -----------------------------
        async for message in websocket:

            if not ROOM["is_active"]:
                break

            data = json.loads(message)

            if data["type"] != "chat":
                continue

            user_text = data["message"].strip()
            player_data = ROOM["players"][websocket]

            # HINT
            if user_text.upper() == "[HINT]":
                if not await handle_hint(websocket, player_name, player_data):
                    break
                continue

            # STANDARD ATTEMPT
            player_data["attempts"] += 1

            if player_data["attempts"] > ROOM["tries"]:
                await send(websocket, {
                    "type": "game_over",
                    "result": "loss",
                    "message": "MAXIMUM ATTEMPTS EXCEEDED. YOU HAVE BEEN CAUGHT."
                })
                break

            await send_admin_chat(player_name, player_name, user_text)

            # WIN CHECK
            if secret_upper in user_text.upper():
                ROOM["is_active"] = False

                await broadcast({
                    "type": "game_over",
                    "result": "global_win",
                    "winner": player_name,
                    "message": f"SYSTEM HACKED BY {player_name}! PASSWORD: {ROOM['secret']}"
                })

                return

            # AI RESPONSE
            history.append({"role": "user", "content": user_text})
            await handle_ai_response(websocket, history, player_name)

    finally:
        ROOM["players"].pop(websocket, None)

        if ROOM["is_active"]:
            await broadcast({
                "type": "system_notice",
                "message": f"OPERATIVE '{player_name}' DISCONNECTED."
            })
        
        if ROOM["admin_ws"]:
            names = [p["name"] for p in ROOM["players"].values()]
            await send(ROOM["admin_ws"], {
                "type": "player_update",
                "count": len(ROOM["players"]),
                "names": names
            })


# =========================================================
# SERVER
# =========================================================

async def main_handler(websocket):
    path = getattr(websocket.request, "path", "/")

    if "/admin" in path:
        await handle_admin(websocket)
    else:
        await handle_player(websocket)

def serve_frontend():
    """Starts a separate HTTP server to handle file requests."""
    class QuietHandler(http.server.SimpleHTTPRequestHandler):
        def log_message(self, format, *args): return # Keep logs clean

    # Port 8000 is shared, but we need to handle this carefully.
    # Actually, let's bind the HTTP server to 8000 and the WS to the SAME loop.
    # To keep it simple and foolproof:
    handler = QuietHandler
    with socketserver.TCPServer(("0.0.0.0", 8000), handler) as httpd:
        print("📂 Web Server (HTTP) listening on port 8000")
        httpd.serve_forever()

async def main():
    # Start the HTTP file server in a separate thread
    # Note: We must use separate ports if we can't multiplex.
    # SHIFT: Let's use 8000 for everything by using the most compatible method.
    
    print("🚀 Starting Agent Breaker Services...")
    
    # Start HTTP server in background
    thread = threading.Thread(target=serve_frontend, daemon=True)
    thread.start()

    # WebSocket server on the SAME port won't work easily in Python without complex 
    # multiplexing. Let's move WebSockets to 8001 internally and 
    # update your docker-compose to map it.
    
    async with websockets.serve(main_handler, "0.0.0.0", 8001):
        print("🔌 WebSocket (WS) listening on port 8001")
        await asyncio.Future()


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