import contextlib
import logging
import multiprocessing
import os
import random
import signal
import socket
import string
import struct
import subprocess
import sys
import tempfile
import traceback

from libosue import *
from libosue import CSnippet as C
from libosue import FunctionPrototype as FP


################################################################################
# INITIALIZE THE EXAM OBJECT
################################################################################

# These functions are not allowed.
not_allowed = [
    "DEMO_setup_connection",
    "DEMO_server_request_handler",
    "DEMO_execute_command",
]

prototypes = {
    "parse_arguments": FP("void")("int argc", "char **argv", "struct arguments *args"),
    "DEMO_setup_connection": FP("int")("const char *port_str"),
    "DEMO_server_request_handler": FP("void")("int sockfd", "const char *command", "int flags"),
    "DEMO_execute_command": FP("FILE *")("const char *command", "const char *argument"),
}

PORT = 692

exam.initialize(
    "./server",
    server_port=7220,
    not_allowed=not_allowed + exam.DEFAULT_RESTRICTED_FUNCTIONS,
    call_limit=200,
    prototypes=prototypes,
    report_file="report.txt",
    result_file="result.txt",
    debug_file="libosue.dbg",
    make_targets=["clean", "all"],
    initialize_stack=["setup_connection", "server_request_handler", "execute_command"],
)


class FunctionAborted(EvalError):
    """Raised when a function should be aborted."""

    pass


# Kill programs on port 692.
# This is slow, so we only do it once every deliver call and not for every session.
utils.free_port(PORT)


################################################################################
# TASK 1: SETUP CONNECTION
################################################################################


def bind_cb(s, args):
    """Call setsockopt(...) before binding."""
    reuse = Value.from_type("char")
    reuse.value = 1
    s.c.setsockopt(args["sockfd"], socket.SOL_SOCKET, socket.SO_REUSEADDR, reuse.address, 1)


with exam.task("Task 1: setup connection", 10) as task:
    with task.subtask("initialize server socket", 8) as st:
        with exam.session([]) as s:
            # Setup tracing for hints.
            s.c.socket.tracing = True
            s.c.bind.tracing = True
            s.c.listen.tracing = True

            # Make sure setsockopt(...) is called before bind().
            s.c.bind.call_cb = bind_cb

            # Call the student's setup_connection connection function
            sockfd = s.c.setup_connection(str(PORT).encode())

            # Check if the server socket works and creates valid connection sockets.
            utils.check_server_socket(s, sockfd, PORT)

    # Give the students some hints for their implementation
    if st.points != 8:
        if "socket" not in s.trace:
            exam.log("[INFO] socket() was never called")
        if "bind" not in s.trace:
            exam.log("[INFO] bind() was never called")
        if "listen" not in s.trace:
            exam.log("[INFO] listen() was never called")

        exam.log("The following subtasks are skipped because of errors:")
        exam.log("* error handling (2 points)")
        task.skip_subtasks()

    with task.subtask("error handling", 2) as st:
        with st.subtask("bind()", 1):
            with exam.session([]) as s:
                s.c.listen.tracing = True
                s.c.bind.override_return = C("((int) -1)")

                with exam.raises(ExitError, "program did not exit after bind() failed") as excinfo:
                    sockfd = s.c.setup_connection(str(PORT).encode())
                assert excinfo.value.exitcode == 1, "wrong exit code"
                assert "listen" not in s.trace, "listen() called after bind() failed"

        with st.subtask("listen()", 1):
            with exam.session([]) as s:
                s.c.listen.override_return = C("((int) -1)")

                with exam.raises(ExitError, "program did not exit after listen() failed") as excinfo:
                    sockfd = s.c.setup_connection(str(PORT).encode())
                assert excinfo.value.exitcode == 1, "wrong exit code"


with exam.task("Task 2: handle requests", 20) as task:
    with task.subtask("read argument and execute command", 7) as st:
        # Check if the program correctly accepts a connection and reads the argument
        # from the connection socket. Execution stops when task 3 is called.
        def accept_return_cb(s, rv):
            rv = rv.convert()
            if rv > 0:
                s["connfd"] = rv

        def read_call_cb(s, args):
            connfd = args.get("sockfd") or args.get("fd")
            s["read"] = connfd.convert() == s["connfd"]
            s["buffer"] = args.get("buf")

        def read_return_cb(s, rv):
            if s["read"]:
                n = rv.convert()
                s["result"] += s["buffer"].convert(shape=n, cast="char *")

        def execute_command_cb(s, args):
            s["ec_called"] = True
            try:
                s["command"] = args["command"].convert()
                s["argument"] = args["argument"].convert()
            except InvalidMemoryError:
                pass
            raise FunctionAborted("This is where execution ends.")

        with exam.session([]) as s:
            s.c.accept.return_cb = accept_return_cb
            s.c.read.call_cb = read_call_cb
            s.c.read.return_cb = read_return_cb
            s.c.recv.call_cb = read_call_cb
            s.c.recv.return_cb = read_return_cb

            s.c.execute_command.call_cb = execute_command_cb
            s.c.DEMO_execute_command.call_cb = execute_command_cb
            s.c.DEMO_execute_command.allowed = True

            # Setup the tracing variables.
            s["connfd"] = None
            s["result"] = b""
            s["ec_called"] = False
            s["command"] = None
            s["argument"] = None

            argument = "".join(random.choices(string.ascii_letters, k=24)).encode()
            sockfd = s.c.DEMO_setup_connection(str(PORT).encode())
            with utils.SingleShotClient(PORT, argument) as client:
                with contextlib.suppress(FunctionAborted):
                    s.c.server_request_handler(sockfd, b"./osue-tool")
            assert s["result"], "no argument read from client"
            assert s["result"] == argument, "read wrong argument from client"
            assert s["ec_called"], "execute_command() not called"
            assert s["command"] == b"./osue-tool", "invalid command passt to execute_command()"
            assert (
                s["argument"] == argument
            ), "Invalid argument passed to execute_command()\nDid you terminate it?"

    with task.subtask("response", 8) as st:
        # Execute the full function and check if the client receives the correct result.
        def execute_command_cb(s, args):
            with s.c.fork.save():
                s.c.fork.allowed = True
                return s.c.DEMO_execute_command(args["command"], args["argument"])

        with st.subtask("single line output", 4) as pt:
            with exam.session([]) as s:
                s.c.execute_command.call_cb = execute_command_cb
                s.c.DEMO_execute_command.call_cb = execute_command_cb
                s.c.DEMO_execute_command.allowed = True

                sockfd = s.c.DEMO_setup_connection(str(PORT).encode())

                argument = "".join(random.choices(string.ascii_letters, k=24)).encode()
                expected = argument + b"\n"
                with utils.SingleShotClient(PORT, argument) as client:
                    s.c.server_request_handler(sockfd, b"./osue-tool")
                assert client.response, "didn't receive a response for single-line output"
                assert client.response == expected, "received wrong response for single-line output"

        with st.subtask("multi line output", 4) as pt:
            with exam.session([]) as s:
                s.c.execute_command.call_cb = execute_command_cb
                s.c.DEMO_execute_command.call_cb = execute_command_cb
                s.c.DEMO_execute_command.allowed = True

                sockfd = s.c.DEMO_setup_connection(str(PORT).encode())

                argument = "".join(random.choices(string.ascii_letters, k=128)).encode()
                expected = b"\n".join(argument[i : i + 32] for i in range(0, len(argument), 32)) + b"\n"
                with utils.SingleShotClient(PORT, argument) as client:
                    s.c.server_request_handler(sockfd, b"./osue-tool")
                assert client.response, "didn't receive a response for multi-line output"
                assert client.response == expected, "received wrong response for multi-line output"

    if task.points != 15:
        exam.log("The following subtasks are skipped because of errors:")
        exam.log("* repeated write/send calls (5 points)")
        task.skip_subtasks()

    with task.subtask("repeated write/send calls", 5) as st:

        def accept_return_cb(s, rv):
            s["connfd"] = rv.convert()

        def write_cb(s, args):
            fd = (args.get("fd") or args.get("sockfd")).convert()

            if fd == s["connfd"]:
                n = args.get("count") or args.get("len")
                buf = args.get("buf")
                s["messages"].append(buf.convert(shape=n.convert(), cast="char *"))
                return 1

        with exam.session([]) as s:
            s.c.execute_command.call_cb = execute_command_cb
            s.c.DEMO_execute_command.call_cb = execute_command_cb
            s.c.DEMO_execute_command.allowed = True

            s.c.accept.return_cb = accept_return_cb
            s.c.write.call_cb = write_cb
            s.c.send.call_cb = write_cb

            s["connfd"] = None
            s["messages"] = []

            argument = "".join(random.choices(string.ascii_letters, k=24)).encode()
            expected = argument + b"\n"
            sockfd = s.c.DEMO_setup_connection(str(PORT).encode())
            with utils.SingleShotClient(PORT, argument) as client:
                s.c.server_request_handler(sockfd, b"./osue-tool")

            assert len(s["messages"]) > 1, "write/send not called multiple times"
            message = b"".join(m[:1] for m in s["messages"])
            assert message == expected, "received wrong response"


with exam.task("Task 3: execute command", 20) as task:
    with task.subtask("fork/exec pattern", 5) as st:
        pattern = utils.ForkNode(parent=None, child=utils.ExecNode("./osue-tool"))
        with utils.ForkExecTraces(pattern) as tracegen:
            for fork_exec_run in tracegen:
                with exam.session() as s:
                    with fork_exec_run(s):
                        s.c.execute_command(b"./osue-tool", b"just a random string")

        assert tracegen.correct, tracegen.default_message

    if task.points != 5:
        exam.log("The following subtasks are skipped because of errors:")
        exam.log("* wait for child process to finish  (5 points)")
        exam.log("* result                            (10 points)")
        task.skip_subtasks()

    with task.subtask("wait for child process to finish", 5) as st:
        # Correctly wait for the child process to finish.
        # Follow the parent process and trace the return value of fork() and wait().
        def fork_return_cb(s, rv):
            s["fork_result"] = rv.convert()

        def wait_return_cb(s, rv):
            s["wait_result"] = rv.convert()

        with exam.session([]) as s:
            s.c.fork.allowed = True
            s.c.fork.return_cb = fork_return_cb
            s.c.wait.return_cb = wait_return_cb
            s.c.waitpid.return_cb = wait_return_cb
            s.c.execute_command(b"./osue-tool", b"just a random string")

            assert "wait_result" in s, "wait never called"
            assert s.get("wait_result") == s.get("fork_result"), "program did not wait for child process"

    if task.points != 10:
        exam.log("The following subtasks are skipped because of errors:")
        exam.log("* result (10 points)")
        task.skip_subtasks()

    with task.subtask("result", 10) as st:
        # Handle a single request and check the result.
        with exam.session([]) as s:
            s.c.fork.allowed = True

            argument = "".join(random.choices(string.ascii_letters, k=24)).encode()
            expected = argument + b"\n"

            fp = s.c.execute_command(b"./osue-tool", argument)
            assert fp.intaddr(), "returned file pointer not valid"

            bufsize = 8192
            try:
                buf = s.c.malloc(bufsize, cast="char *")
                n = s.c.fread(buf, 1, bufsize, fp)
            except FunctionKilledError:
                st.fail("reading from returned file pointer never terminated")
            except SignalError:
                st.fail("reading from returned file pointer failed")

            result = buf.convert(shape=n.convert())
            assert result == expected, "invalid response"