import asyncio
import argparse
import logging
import sys

from server.server import GameServer
from server.config import ServerConfig


async def _run_tasks_with_optional_timeout(tasks, timeout_s=None):
    """Await tasks, optionally with a timeout to cancel after specified seconds."""
    if not timeout_s:
        await asyncio.gather(*tasks)
        return
    try:
        await asyncio.wait_for(asyncio.gather(*tasks), timeout=float(timeout_s))
    except asyncio.TimeoutError:
        logging.info("Timeout reached (%s seconds); stopping server tasks...", timeout_s)
        for t in tasks:
            t.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)


async def run_in_memory(args):
    from server.transport import InMemoryTransport
    cfg = ServerConfig(
        width=args.width,
        height=args.height,
        tick_rate=args.tick_rate,
        wrap_edges=args.wrap_edges,
        apples_per_snake=args.apples_per_snake,
        apples_cap=args.apples_cap,
        players_max=args.players_max,
    )
    server = GameServer(transport=InMemoryTransport(lambda d, p: server.on_datagram(d, p)), config=cfg)
    tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
    await _run_tasks_with_optional_timeout(tasks, args.run_seconds)


async def run_quic(args):
    import os
    from server.quic_transport import QuicWebTransportServer
    from server.utils import combine_cert_chain

    cfg = ServerConfig(
        width=args.width,
        height=args.height,
        tick_rate=args.tick_rate,
        wrap_edges=args.wrap_edges,
        apples_per_snake=args.apples_per_snake,
        apples_cap=args.apples_cap,
        players_max=args.players_max,
    )

    # Handle certificate chain
    # Check if cert file contains multiple certificates (like Let's Encrypt fullchain.pem)
    with open(args.cert, 'r') as f:
        cert_content = f.read()

    from server.utils import split_pem_certificates
    certs_in_file = split_pem_certificates(cert_content)

    cert_file = args.cert
    temp_cert_file = None

    if args.cert_chain:
        # Combine cert + chain files into single file
        temp_cert_file = combine_cert_chain(args.cert, args.cert_chain)
        cert_file = temp_cert_file
        logging.info("Combined certificate with %d chain file(s)", len(args.cert_chain))
    elif len(certs_in_file) > 1:
        # Certificate file contains full chain - reformat it for aioquic
        temp_cert_file = combine_cert_chain(args.cert, [])
        cert_file = temp_cert_file
        logging.info("Reformatted certificate chain from fullchain file (%d certs)", len(certs_in_file))

    server = GameServer(transport=QuicWebTransportServer(args.quic_host, args.quic_port, cert_file, args.key, lambda d, p: server.on_datagram(d, p)), config=cfg)
    logging.info("QUIC server: %s:%d", args.quic_host, args.quic_port)
    try:
        tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
        await _run_tasks_with_optional_timeout(tasks, args.run_seconds)
    finally:
        if temp_cert_file:
            os.unlink(temp_cert_file)  # Clean up temp combined cert file


async def run_webtransport(args):
    import json
    import os
    from server.webtransport_server import WebTransportServer
    from server.static_server import start_https_static
    from server.utils import get_cert_sha256_hash, combine_cert_chain, extract_server_cert

    cfg = ServerConfig(
        width=args.width,
        height=args.height,
        tick_rate=args.tick_rate,
        wrap_edges=args.wrap_edges,
        apples_per_snake=args.apples_per_snake,
        apples_cap=args.apples_cap,
        players_max=args.players_max,
    )

    # Handle certificate chain
    # Check if cert file contains multiple certificates (like Let's Encrypt fullchain.pem)
    with open(args.cert, 'r') as f:
        cert_content = f.read()

    from server.utils import split_pem_certificates
    certs_in_file = split_pem_certificates(cert_content)

    cert_file = args.cert
    temp_cert_file = None

    if args.cert_chain:
        # Combine cert + chain files into single file
        temp_cert_file = combine_cert_chain(args.cert, args.cert_chain)
        cert_file = temp_cert_file
        logging.info("Combined certificate with %d chain file(s)", len(args.cert_chain))
    elif len(certs_in_file) > 1:
        # Certificate file contains full chain - reformat it for aioquic
        temp_cert_file = combine_cert_chain(args.cert, [])
        cert_file = temp_cert_file
        logging.info("Reformatted certificate chain from fullchain file (%d certs)", len(certs_in_file))

    # Calculate certificate hash for WebTransport client (only hash server cert, not chain)
    server_cert_file = extract_server_cert(args.cert)
    cert_hash = get_cert_sha256_hash(server_cert_file)
    if server_cert_file != args.cert:
        os.unlink(server_cert_file)  # Clean up temp file

    # Prepare cert-hash.json content
    cert_hash_json = json.dumps({
        "sha256": cert_hash,
        "wtUrl": f"https://{args.wt_host}:{args.wt_port}/",
        "wtPort": args.wt_port
    })

    # Optional static HTTPS server for client assets
    httpd = None
    if args.static:
        httpd, _t = start_https_static(
            args.static_host,
            args.static_port,
            cert_file,
            args.key,
            args.static_root,
            cert_hash_json=cert_hash_json
        )
        print(f"HTTPS static server: https://{args.static_host}:{args.static_port}/ serving '{args.static_root}'")
        print(f"Certificate SHA-256: {cert_hash}")

    server = GameServer(transport=WebTransportServer(args.wt_host, args.wt_port, cert_file, args.key, lambda d, p: server.on_datagram(d, p)), config=cfg)
    print(f"WebTransport server: https://{args.wt_host}:{args.wt_port}/ (HTTP/3)")
    try:
        tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
        await _run_tasks_with_optional_timeout(tasks, args.run_seconds)
    finally:
        if httpd is not None:
            httpd.shutdown()
        if temp_cert_file:
            os.unlink(temp_cert_file)  # Clean up temp combined cert file


async def run_net(args):
    import os
    from server.webtransport_server import WebTransportServer
    from server.quic_transport import QuicWebTransportServer
    from server.multi_transport import MultiTransport
    from server.utils import combine_cert_chain

    cfg = ServerConfig(
        width=args.width,
        height=args.height,
        tick_rate=args.tick_rate,
        wrap_edges=args.wrap_edges,
        apples_per_snake=args.apples_per_snake,
        apples_cap=args.apples_cap,
        players_max=args.players_max,
    )

    # Handle certificate chain
    # Check if cert file contains multiple certificates (like Let's Encrypt fullchain.pem)
    with open(args.cert, 'r') as f:
        cert_content = f.read()

    from server.utils import split_pem_certificates
    certs_in_file = split_pem_certificates(cert_content)

    cert_file = args.cert
    temp_cert_file = None

    if args.cert_chain:
        # Combine cert + chain files into single file
        temp_cert_file = combine_cert_chain(args.cert, args.cert_chain)
        cert_file = temp_cert_file
        logging.info("Combined certificate with %d chain file(s)", len(args.cert_chain))
    elif len(certs_in_file) > 1:
        # Certificate file contains full chain - reformat it for aioquic
        temp_cert_file = combine_cert_chain(args.cert, [])
        cert_file = temp_cert_file
        logging.info("Reformatted certificate chain from fullchain file (%d certs)", len(certs_in_file))

    server: GameServer
    wt = WebTransportServer(args.wt_host, args.wt_port, cert_file, args.key, lambda d, p: server.on_datagram(d, p))
    qu = QuicWebTransportServer(args.quic_host, args.quic_port, cert_file, args.key, lambda d, p: server.on_datagram(d, p))
    m = MultiTransport(wt, qu)
    server = GameServer(transport=m, config=cfg)
    logging.info("WebTransport server: %s:%d", args.wt_host, args.wt_port)
    logging.info("QUIC server: %s:%d", args.quic_host, args.quic_port)
    try:
        tasks = [asyncio.create_task(m.run()), asyncio.create_task(server.tick_loop())]
        await _run_tasks_with_optional_timeout(tasks, args.run_seconds)
    finally:
        if temp_cert_file:
            os.unlink(temp_cert_file)  # Clean up temp combined cert file


def parse_args():
    parser = argparse.ArgumentParser(
        description="Real-time multiplayer Snake game server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run.py                                    # WebTransport mode with defaults
  python run.py --mode mem                         # In-memory mode (testing)
  python run.py --mode quic --quic-port 5000       # QUIC mode on custom port
  python run.py --cert my.pem --key my.key         # Custom TLS certificate
  python run.py --no-static                        # Disable static HTTPS server
  python run.py --width 80 --height 60             # Custom field size
        """
    )

    # Transport mode
    parser.add_argument("--mode", choices=["mem", "wt", "quic", "net"], default="wt",
                        help="Transport mode: mem (in-memory), wt (WebTransport/HTTP3), quic (QUIC datagrams), net (both wt+quic). Default: wt")

    # TLS certificate/key
    parser.add_argument("--cert", default="cert.pem",
                        help="TLS certificate file path (may contain full chain). Default: cert.pem")
    parser.add_argument("--key", default="key.pem",
                        help="TLS private key file path. Default: key.pem")
    parser.add_argument("--cert-chain", action="append", dest="cert_chain",
                        help="Intermediate certificate file (can be used multiple times for long chains)")

    # WebTransport server settings
    parser.add_argument("--wt-host", default="0.0.0.0",
                        help="WebTransport server host. Default: 0.0.0.0")
    parser.add_argument("--wt-port", type=int, default=4433,
                        help="WebTransport server port. Default: 4433")

    # QUIC server settings
    parser.add_argument("--quic-host", default="0.0.0.0",
                        help="QUIC server host. Default: 0.0.0.0")
    parser.add_argument("--quic-port", type=int, default=4433,
                        help="QUIC server port. Default: 4433 (or 4443 in net mode)")

    # Static HTTPS server settings
    parser.add_argument("--static", dest="static", action="store_true", default=True,
                        help="Enable static HTTPS server for client files (default in wt mode)")
    parser.add_argument("--no-static", dest="static", action="store_false",
                        help="Disable static HTTPS server")
    parser.add_argument("--static-host", default=None,
                        help="Static server host. Default: same as wt-host")
    parser.add_argument("--static-port", type=int, default=8443,
                        help="Static server port. Default: 8443")
    parser.add_argument("--static-root", default="client",
                        help="Static files directory. Default: client")

    # Game configuration
    parser.add_argument("--width", type=int, default=60,
                        help="Field width (3-255). Default: 60")
    parser.add_argument("--height", type=int, default=40,
                        help="Field height (3-255). Default: 40")
    parser.add_argument("--tick-rate", type=int, default=10,
                        help="Server tick rate in TPS (5-30). Default: 10")
    parser.add_argument("--wrap-edges", action="store_true", default=False,
                        help="Enable edge wrapping (default: disabled)")
    parser.add_argument("--apples-per-snake", type=int, default=1,
                        help="Apples per connected snake (1-12). Default: 1")
    parser.add_argument("--apples-cap", type=int, default=255,
                        help="Maximum total apples (0-255). Default: 255")
    parser.add_argument("--players-max", type=int, default=32,
                        help="Maximum concurrent players. Default: 32")

    # Logging and testing
    parser.add_argument("--log-level", default="INFO",
                        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
                        help="Logging level. Default: INFO")
    parser.add_argument("--run-seconds", type=float, default=None,
                        help="Optional timeout in seconds for testing")

    args = parser.parse_args()

    # Post-process: static-host defaults to wt-host
    if args.static_host is None:
        args.static_host = args.wt_host

    # In net mode, default quic-port to 4443 if not explicitly set
    if args.mode == "net" and "--quic-port" not in sys.argv:
        args.quic_port = 4443

    # Validate TLS files for modes that need them
    if args.mode in ("wt", "quic", "net"):
        import os
        if not os.path.exists(args.cert):
            parser.error(f"Certificate file not found: {args.cert}")
        if not os.path.exists(args.key):
            parser.error(f"Key file not found: {args.key}")

    return args


if __name__ == "__main__":
    try:
        args = parse_args()

        # Logging setup
        logging.basicConfig(
            level=getattr(logging, args.log_level.upper()),
            format="[%(asctime)s] %(levelname)s: %(message)s"
        )

        # Run appropriate mode
        if args.mode == "wt":
            logging.info("Starting in WebTransport mode")
            asyncio.run(run_webtransport(args))
        elif args.mode == "quic":
            logging.info("Starting in QUIC datagram mode")
            asyncio.run(run_quic(args))
        elif args.mode == "net":
            logging.info("Starting in combined WebTransport+QUIC mode")
            asyncio.run(run_net(args))
        else:  # mem
            logging.info("Starting in in-memory transport mode")
            asyncio.run(run_in_memory(args))
    except KeyboardInterrupt:
        pass