import asyncio
import socket
import ssl
import capnp
import logging

connection_capnp = capnp.load('schema/connection.capnp')
authenticationsystem_capnp = capnp.load('schema/authenticationsystem.capnp')


async def myreader(client, reader):
    while True:
        data = await reader.read(4096)
        client.write(data)


async def mywriter(client, writer):
    while True:
        data = await client.read(4096)
        writer.write(data.tobytes())
        await writer.drain()


async def connect(host, port, user, pw):
    # Setup SSL context
    ctx = ssl.create_default_context(
        ssl.Purpose.SERVER_AUTH
    )
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    # Handle both IPv4 and IPv6 cases
    try:
        reader, writer = await asyncio.open_connection(
            host, port, ssl=ctx, family=socket.AF_INET6
        )
    except Exception:
        reader, writer = await asyncio.open_connection(
            host, port, ssl=ctx, family=socket.AF_INET
        )

    # Start TwoPartyClient using TwoWayPipe (takes no arguments in this mode)
    client = capnp.TwoPartyClient()
    cap = client.bootstrap().cast_as(connection_capnp.Bootstrap)

    # Assemble reader and writer tasks, run in the background
    coroutines = [myreader(client, reader), mywriter(client, writer)]
    asyncio.gather(*coroutines, return_exceptions=True)

    auth = cap.authenticationSystem().authenticationSystem
    req = auth.start_request()
    req.request.mechanism = "PLAIN"
    req.request.initialResponse.initial = "\0" + user + "\0" + pw
    rep_prom = req.send()
    response = await rep_prom.a_wait()
    if response.response.outcome.result == "successful":
        logging.info("Auth completed successfully")
        return cap
    else:
        logging.info("Auth failed")
        return None


async def connect_with_fabfire_initial(host, port, uid):
    # Setup SSL context
    ctx = ssl.create_default_context(
        ssl.Purpose.SERVER_AUTH
    )
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    # Handle both IPv4 and IPv6 cases
    try:
        reader, writer = await asyncio.open_connection(
            host, port, ssl=ctx, family=socket.AF_INET6
        )
    except Exception:
        reader, writer = await asyncio.open_connection(
            host, port, ssl=ctx, family=socket.AF_INET
        )

    # Start TwoPartyClient using TwoWayPipe (takes no arguments in this mode)
    client = capnp.TwoPartyClient()
    cap = client.bootstrap().cast_as(connection_capnp.Bootstrap)

    # Assemble reader and writer tasks, run in the background
    coroutines = [myreader(client, reader), mywriter(client, writer)]
    asyncio.gather(*coroutines, return_exceptions=True)

    auth = cap.authenticationSystem().authenticationSystem
    req = auth.start_request()
    req.request.mechanism = "X-FABFIRE"
    req.request.initialResponse.initial = uid
    rep_prom = req.send()
    response = await rep_prom.a_wait()
    logging.debug(f"got response type: {response.response.which()}")
    if response.response.which() == "challence":
        logging.debug(f"challenge: {response.response.challence}")
        return auth, response.response.challence, cap
    else:
        logging.error(f"Auth failed: {response.response.outcome.result}, additional info: {response.response.outcome.helpText}")
        return None


async def connect_with_fabfire_step(auth, msg):
    req = auth.step_request()
    req.response = msg
    rep_prom = req.send()
    response = await rep_prom.a_wait()
    if response.response.which() == "challence":
        logging.debug(f"challenge: {response.response.challence}")
        return response.response.challence, False  # auth cap, challenge, not done
    elif response.response.outcome.result == "successful":
        logging.info(f"Auth completed successfully! Got additional Data: {response.response.outcome.additionalData.additional}")
        return response.response.outcome.additionalData.additional, True  # dont care, message, we are done
    else:
        logging.error(f"Auth failed: {response.response.outcome.result}, additional info: {response.response.outcome.helpText}")
        return None