"""
CruiseBuddy Server v1
=====================
Real-time group drive tracking with WebSockets.

Requirements:
    pip install aiohttp

Run:
    python server.py

Then open http://localhost:8765 in your browser.
To make it accessible on your network, it binds to 0.0.0.0 by default.
"""

import asyncio
import json
import string
import random
import time
import os
from collections import OrderedDict
from aiohttp import web

# ═══════════════════════════════════════════
# CONFIG
# ═══════════════════════════════════════════
HOST = "0.0.0.0"
PORT = 8765
STATIC_DIR = os.path.dirname(os.path.abspath(__file__))

# ═══════════════════════════════════════════
# STATE
# ═══════════════════════════════════════════
rooms = {}        # code -> Room dict
ws_to_room = {}   # websocket -> room code
ws_to_name = {}   # websocket -> display name


def generate_code():
    """Generate a 5-character room code (no ambiguous chars)."""
    chars = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
    while True:
        code = "".join(random.choices(chars, k=5))
        if code not in rooms:
            return code


def get_room(ws):
    """Get the room a websocket belongs to."""
    code = ws_to_room.get(ws)
    if code and code in rooms:
        return rooms[code]
    return None


def get_participant(room, ws):
    """Get participant dict from room by websocket."""
    for p in room["participants"]:
        if p["ws"] is ws:
            return p
    return None


def participant_list(room):
    """Build a serializable participant list (no ws objects)."""
    return [
        {
            "name": p["name"],
            "role": p["role"],
        }
        for p in room["participants"]
    ]


async def broadcast(room, message, exclude=None):
    """Send message to all participants in a room."""
    data = json.dumps(message)
    for p in room["participants"]:
        if p["ws"] is not exclude:
            try:
                await p["ws"].send_str(data)
            except Exception:
                pass


async def send(ws, message):
    """Send message to a single websocket."""
    try:
        await ws.send_str(json.dumps(message))
    except Exception:
        pass


async def send_error(ws, msg):
    await send(ws, {"type": "error", "msg": msg})


# ═══════════════════════════════════════════
# MESSAGE HANDLERS
# ═══════════════════════════════════════════

async def handle_create_room(ws, data):
    name = data.get("name", "").strip()[:20]
    max_p = data.get("max_participants", 5)
    max_p = max(2, min(10, int(max_p)))

    if not name:
        return await send_error(ws, "Name is required.")

    # Leave any existing room first
    await handle_leave(ws, silent=True)

    code = generate_code()
    room = {
        "code": code,
        "leader_ws": ws,
        "leader_name": name,
        "max_participants": max_p,
        "status": "lobby",  # lobby | active | paused
        "participants": [
            {
                "ws": ws,
                "name": name,
                "role": "leader",
                "lat": None,
                "lng": None,
                "speed": 0,
                "heading": 0,
            }
        ],
        "chat": [],
        "created_at": time.time(),
    }

    rooms[code] = room
    ws_to_room[ws] = code
    ws_to_name[ws] = name

    await send(ws, {
        "type": "room_created",
        "code": code,
        "name": name,
        "max_participants": max_p,
        "participants": participant_list(room),
    })

    print(f"[ROOM] {name} created room {code} (max {max_p})")


async def handle_join_room(ws, data):
    name = data.get("name", "").strip()[:20]
    code = data.get("code", "").strip().upper()

    if not name:
        return await send_error(ws, "Name is required.")
    if not code or len(code) != 5:
        return await send_error(ws, "Invalid room code.")

    room = rooms.get(code)
    if not room:
        return await send_error(ws, "Room not found.")

    if len(room["participants"]) >= room["max_participants"]:
        return await send_error(ws, "Room is full.")

    # Check for duplicate name
    for p in room["participants"]:
        if p["name"].lower() == name.lower():
            return await send_error(ws, "Name already taken in this room.")

    # Leave any existing room first
    await handle_leave(ws, silent=True)

    participant = {
        "ws": ws,
        "name": name,
        "role": "member",
        "lat": None,
        "lng": None,
        "speed": 0,
        "heading": 0,
    }

    room["participants"].append(participant)
    ws_to_room[ws] = code
    ws_to_name[ws] = name

    # Send room state to joiner
    await send(ws, {
        "type": "joined",
        "code": code,
        "role": "member",
        "status": room["status"],
        "leader_name": room["leader_name"],
        "max_participants": room["max_participants"],
        "participants": participant_list(room),
        "chat": room["chat"][-50:],  # last 50 messages
    })

    # Broadcast updated participant list to everyone else
    await broadcast(room, {
        "type": "participant_update",
        "participants": participant_list(room),
        "event": "joined",
        "name": name,
    }, exclude=ws)

    print(f"[JOIN] {name} joined room {code} ({len(room['participants'])}/{room['max_participants']})")


async def handle_start_drive(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")
    if room["leader_ws"] is not ws:
        return await send_error(ws, "Only the leader can start.")
    if room["status"] != "lobby":
        return await send_error(ws, "Drive already started.")

    room["status"] = "active"
    await broadcast(room, {"type": "drive_started"})
    print(f"[DRIVE] Room {room['code']} drive started")


async def handle_pause_drive(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")
    if room["leader_ws"] is not ws:
        return await send_error(ws, "Only the leader can pause.")
    if room["status"] != "active":
        return await send_error(ws, "Drive is not active.")

    room["status"] = "paused"
    await broadcast(room, {"type": "drive_paused"})
    print(f"[DRIVE] Room {room['code']} drive paused")


async def handle_resume_drive(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")
    if room["leader_ws"] is not ws:
        return await send_error(ws, "Only the leader can resume.")
    if room["status"] != "paused":
        return await send_error(ws, "Drive is not paused.")

    room["status"] = "active"
    await broadcast(room, {"type": "drive_resumed"})
    print(f"[DRIVE] Room {room['code']} drive resumed")


async def handle_end_drive(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")
    if room["leader_ws"] is not ws:
        return await send_error(ws, "Only the leader can end the drive.")

    room["status"] = "lobby"
    # Reset all positions
    for p in room["participants"]:
        p["lat"] = None
        p["lng"] = None
        p["speed"] = 0
        p["heading"] = 0

    await broadcast(room, {
        "type": "drive_ended",
        "participants": participant_list(room),
    })
    print(f"[DRIVE] Room {room['code']} drive ended")


async def handle_kick(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")
    if room["leader_ws"] is not ws:
        return await send_error(ws, "Only the leader can kick.")

    target_name = data.get("target", "").strip()
    target_p = None
    for p in room["participants"]:
        if p["name"] == target_name and p["role"] != "leader":
            target_p = p
            break

    if not target_p:
        return await send_error(ws, "Participant not found.")

    target_ws = target_p["ws"]

    # Remove from room
    room["participants"].remove(target_p)
    if target_ws in ws_to_room:
        del ws_to_room[target_ws]
    if target_ws in ws_to_name:
        del ws_to_name[target_ws]

    # Tell the kicked person
    await send(target_ws, {"type": "you_kicked"})

    # Tell everyone else
    await broadcast(room, {
        "type": "participant_update",
        "participants": participant_list(room),
        "event": "kicked",
        "name": target_name,
    })

    print(f"[KICK] {target_name} kicked from room {room['code']}")


async def handle_position(ws, data):
    room = get_room(ws)
    if not room:
        return
    if room["status"] != "active":
        return

    p = get_participant(room, ws)
    if not p:
        return

    p["lat"] = data.get("lat")
    p["lng"] = data.get("lng")
    p["speed"] = data.get("speed", 0)
    p["heading"] = data.get("heading", 0)

    if p["role"] == "leader":
        # Leader position → send to all followers
        leader_data = {
            "type": "leader_position",
            "lat": p["lat"],
            "lng": p["lng"],
            "speed": p["speed"],
            "heading": p["heading"],
        }
        await broadcast(room, leader_data, exclude=ws)
    else:
        # Follower position → send to leader only
        if room["leader_ws"]:
            follower_data = {
                "type": "follower_position",
                "name": p["name"],
                "lat": p["lat"],
                "lng": p["lng"],
                "speed": p["speed"],
                "heading": p["heading"],
            }
            await send(room["leader_ws"], follower_data)


async def handle_chat(ws, data):
    room = get_room(ws)
    if not room:
        return await send_error(ws, "Not in a room.")

    text = data.get("text", "").strip()[:500]
    if not text:
        return

    p = get_participant(room, ws)
    if not p:
        return

    msg = {
        "name": p["name"],
        "text": text,
        "time": int(time.time() * 1000),
        "role": p["role"],
    }

    room["chat"].append(msg)
    # Keep last 200 messages
    if len(room["chat"]) > 200:
        room["chat"] = room["chat"][-200:]

    await broadcast(room, {
        "type": "chat_msg",
        **msg,
    })


async def handle_leave(ws, data=None, silent=False):
    code = ws_to_room.get(ws)
    if not code:
        return

    room = rooms.get(code)
    if not room:
        ws_to_room.pop(ws, None)
        ws_to_name.pop(ws, None)
        return

    p = get_participant(room, ws)
    if not p:
        ws_to_room.pop(ws, None)
        ws_to_name.pop(ws, None)
        return

    name = p["name"]
    is_leader = p["role"] == "leader"

    # Remove participant
    room["participants"] = [x for x in room["participants"] if x["ws"] is not ws]
    ws_to_room.pop(ws, None)
    ws_to_name.pop(ws, None)

    if is_leader or len(room["participants"]) == 0:
        # Leader left or room empty → close the room
        if not silent:
            await broadcast(room, {"type": "room_closed", "reason": "Leader left the drive."})

        # Clean up all remaining participants
        for rp in room["participants"]:
            ws_to_room.pop(rp["ws"], None)
            ws_to_name.pop(rp["ws"], None)

        del rooms[code]
        print(f"[ROOM] Room {code} closed (leader left or empty)")
    else:
        # Regular member left
        if not silent:
            await broadcast(room, {
                "type": "participant_update",
                "participants": participant_list(room),
                "event": "left",
                "name": name,
            })
        print(f"[LEAVE] {name} left room {code}")


# ═══════════════════════════════════════════
# WEBSOCKET HANDLER
# ═══════════════════════════════════════════

MESSAGE_HANDLERS = {
    "create_room": handle_create_room,
    "join_room": handle_join_room,
    "start_drive": handle_start_drive,
    "pause_drive": handle_pause_drive,
    "resume_drive": handle_resume_drive,
    "end_drive": handle_end_drive,
    "kick": handle_kick,
    "position": handle_position,
    "chat": handle_chat,
    "leave": handle_leave,
}


async def websocket_handler(request):
    ws = web.WebSocketResponse(heartbeat=30)
    await ws.prepare(request)

    remote = request.remote or "unknown"
    print(f"[WS] Connected: {remote}")

    try:
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    msg_type = data.get("type", "")
                    handler = MESSAGE_HANDLERS.get(msg_type)
                    if handler:
                        await handler(ws, data)
                    else:
                        await send_error(ws, f"Unknown message type: {msg_type}")
                except json.JSONDecodeError:
                    await send_error(ws, "Invalid JSON.")
                except Exception as e:
                    print(f"[ERROR] {e}")
                    await send_error(ws, "Server error.")
            elif msg.type == web.WSMsgType.ERROR:
                print(f"[WS] Error: {ws.exception()}")
    finally:
        # Clean up on disconnect
        await handle_leave(ws)
        print(f"[WS] Disconnected: {remote}")

    return ws


# ═══════════════════════════════════════════
# HTTP HANDLERS (serve frontend)
# ═══════════════════════════════════════════

async def handle_index(request):
    index_path = os.path.join(STATIC_DIR, "index.html")
    if os.path.exists(index_path):
        return web.FileResponse(index_path)
    return web.Response(text="index.html not found. Place it next to server.py.", status=404)


async def handle_health(request):
    """Simple health/status endpoint."""
    active_rooms = len(rooms)
    total_users = sum(len(r["participants"]) for r in rooms.values())
    return web.json_response({
        "status": "ok",
        "rooms": active_rooms,
        "users": total_users,
        "uptime": time.time(),
    })


# ═══════════════════════════════════════════
# APP SETUP
# ═══════════════════════════════════════════

def create_app():
    app = web.Application()
    app.router.add_get("/", handle_index)
    app.router.add_get("/ws", websocket_handler)
    app.router.add_get("/health", handle_health)
    return app


if __name__ == "__main__":
    print(f"""
╔══════════════════════════════════════════╗
║          CruiseBuddy Server v1           ║
╠══════════════════════════════════════════╣
║  HTTP:  http://localhost:{PORT}             ║
║  WS:    ws://localhost:{PORT}/ws            ║
║                                          ║
║  Listening on {HOST}:{PORT}               ║
║  Press Ctrl+C to stop                    ║
╚══════════════════════════════════════════╝
""")
    app = create_app()
    web.run_app(app, host=HOST, port=PORT, print=None)