"""
WebSocket lifecycle and handler functions.

This module contains the core WebSocket connection handlers, lifecycle management
functions, and background tasks for the BaoLife game server.

Functions:
    - start: Handles WebSocket connection initialization and game loading
    - shutdown: Handles WebSocket disconnection and cleanup
    - error: Sends error messages to clients
    - every_minute: Background task for periodic game iteration
    - handler: Main WebSocket connection handler
    - initialize_dummy_user: Creates a dummy user for server initialization
"""

import asyncio
import datetime
import json
import traceback
import websockets
from websockets import exceptions as websockets_exceptions

from functions import (
    loadGameAsync,
    saveGameAsync,
    playerClass,
    getOccupations,
    insertGame,
    connect
)
from config import config
from server.websocket_registry import USERS
from server.websocket_messaging import sendUserInfo, sendDict


async def start(websocket):
    """
    Initialize a WebSocket connection and load/create player game.

    Handles:
    - Loading existing games from database
    - Creating new games for first-time players
    - Reconnecting disconnected players
    - Initializing game controller and connection state
    - Sending initial player data to client
    - Checking daily login rewards

    Args:
        websocket: WebSocket connection with userID attribute

    Returns:
        None - Starts consumer and producer handlers via asyncio.gather
    """
    # Import from app module
    import app
    consumer_handler = app.consumer_handler
    producer_handler = app.producer_handler
    playerRecords = app.playerRecords

    connected = {websocket}
    print('starting connection...')
    player = False

    cached_player = playerRecords.get(websocket.userID)
    if cached_player is None: #player records are currently loaded games
        print('attempting load for '+websocket.userID)
        player = await loadGameAsync(websocket.userID)
        if (player):
            if (player.gameSpeed == config.SPEED_QUESTION_PAUSE):
                player.gameSpeed = config.SPEED_DEFAULT
            player.controller = 'active'
            player.connection = 'connected'
            connect(player)
            playerRecords.set(websocket.userID, player)
            print('loaded game')
        if (not player):
            print('no loaded games, no saved game -- creating new game')
            player = playerClass()
            player.occupations = getOccupations()
            player.id = websocket.userID

            connect(player)
            insertGame(player)
            playerRecords.set(websocket.userID, player)
    else:
        player = cached_player
        print('reconnected '+player.c.firstname + ' ' + player.c.lastname + ' age ' + str(player.c.ageYears))
        if (player.gameSpeed == config.SPEED_QUESTION_PAUSE):
            player.gameSpeed = config.SPEED_DEFAULT
        player.controller = 'active'
        connect(player)
    await sendUserInfo(player,websocket)
    print(f"USERS: {USERS.count()} Players: {playerRecords.size()}")

    # Check daily login rewards on connect
    from retention.daily_rewards import handle_daily_login_check
    def send_to_client_wrapper(player_id, message):
        asyncio.create_task(sendDict(websocket, message))
    handle_daily_login_check(player.id, send_to_client_wrapper)

    await asyncio.gather(
        consumer_handler(websocket),
        producer_handler(websocket),
    )


async def shutdown(websocket):
    """
    Handle WebSocket disconnection and cleanup.

    Performs:
    - Resets offline tracking
    - Updates player connection state
    - Saves game to database
    - Removes from active users registry

    Args:
        websocket: WebSocket connection with userID attribute

    Returns:
        bool: True if cleanup successful, None otherwise
    """
    # Import playerRecords from app module
    import app
    playerRecords = app.playerRecords

    print(websockets_exceptions.ConnectionClosed)
    print("Client disconnected.  Do cleanup ")
    player = playerRecords.get(websocket.userID)
    if player:
        player.offlineStats.minutesOffline = 0
        player.connection = 'disconnected'
        player.controller = 'inactive'
        await saveGameAsync(player)

        USERS.remove(websocket)
        return True


async def error(websocket, message):
    """
    Send an error message to the client.

    Args:
        websocket: WebSocket connection to send error to
        message: Error message string

    Returns:
        None
    """
    event = {
        "type": "error",
        "message": message,
    }
    await websocket.send(json.dumps(event))


async def every_minute():
    """
    Background task that runs every minute to iterate offline games.

    This task:
    - Checks for disconnected games that need updates
    - Iterates game state for offline players (1 game-minute per real-world minute)
    - Logs player cache statistics

    Runs indefinitely until server shutdown.

    Returns:
        None
    """
    # Import from app module
    import app
    playerRecords = app.playerRecords
    from app import iterateGames

    while True:
        print('every minute, checking for offline games')
        await iterateGames()

        # Log cache stats
        stats = playerRecords.get_stats()
        print(f"PlayerCache: {stats['size']}/{stats['max_size']} players "
              f"({stats['connected']} connected, {stats['disconnected']} disconnected), "
              f"{stats['memory_mb']:.1f}MB")

        await asyncio.sleep(60)  # Run every 60 seconds (1 real-world minute)


async def handler(websocket):
    """
    Main WebSocket connection handler.

    This is the entry point for all WebSocket connections. It:
    - Receives and validates the initial "init" message
    - Extracts userID from the init message
    - Registers the connection in the USERS registry
    - Delegates to start() for game initialization
    - Handles connection errors and cleanup

    Args:
        websocket: WebSocket connection object

    Returns:
        bool: False on error or completion
    """
    # Receive and parse the "init" event from the UI.
    try:
        print(datetime.datetime.now())
        message = await websocket.recv()
        event = json.loads(message)
        if event["type"] == "init":
            websocket.userID = event["userID"]
            print("Client connected with userID " + websocket.userID)
            USERS.add(websocket)  # Upsert behavior
            await start(websocket)
        else:
            websocket.userID = event["userID"]
            print("Client connected with userID " + websocket.userID)
            USERS.add(websocket)  # Upsert behavior
            await start(websocket)
            await error(websocket, "First message must be 'init'")
            await websocket.close()
            return False
        await websocket.close()
    except websockets_exceptions.ConnectionClosed as err:
        print("handler: Client disconnected.  Do cleanup" + str(err))
        await shutdown(websocket)
        return False


async def initialize_dummy_user():
    """
    Create a dummy user for server initialization testing.

    This function connects to the local WebSocket server and sends a test
    message to ensure the server is properly initialized before accepting
    real connections.

    Returns:
        bool: False after completion
    """
    uri = "ws://localhost:8001"  # Adjust the URL as per your setup
    async with websockets.connect(uri) as websocket:
        # Send dummy init message or any other required setup
        dummy_message = {"type": "test", "userID": "DUMMY_USER_ID"}
        try:
            await websocket.send(json.dumps(dummy_message))
            # You might want to do more here or just connect and disconnect
            await websocket.close()
        except websockets_exceptions.ConnectionClosed as err:
            print("initialize_dummy_user: Client disconnected.  Do cleanup" + str(err))
            await shutdown(websocket)
    return False