from typing import Dict, List, Optional, Union
import sys
import json
import contextlib
import io
import traceback
import uuid
import ast

STDIN = sys.stdin
STDOUT = sys.stdout
STDERR = sys.stderr
USER_GLOBALS = {}


def send_message(msg: str):
    length_msg = len(msg)
    STDOUT.buffer.write(f"Content-Length: {length_msg}\r\n\r\n{msg}".encode(encoding="utf-8"))
    STDOUT.buffer.flush()


def print_log(msg: str):
    send_message(json.dumps({"jsonrpc": "2.0", "method": "log", "params": msg}))


def send_response(response: str, response_id: int):
    send_message(json.dumps({"jsonrpc": "2.0", "id": response_id, "result": response}))


def send_request(params: Optional[Union[List, Dict]] = None):
    request_id = uuid.uuid4().hex
    if params is None:
        send_message(json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input"}))
    else:
        send_message(
            json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input", "params": params})
        )
    return request_id


original_input = input


def custom_input(prompt=""):
    try:
        send_request({"prompt": prompt})
        headers = get_headers()
        content_length = int(headers.get("Content-Length", 0))

        if content_length:
            message_text = STDIN.read(content_length)
            message_json = json.loads(message_text)
            our_user_input = message_json["result"]["userInput"]
            return our_user_input
    except Exception:
        print_log(traceback.format_exc())


# Set input to our custom input
USER_GLOBALS["input"] = custom_input
input = custom_input


def handle_response(request_id):
    while not STDIN.closed:
        try:
            headers = get_headers()
            content_length = int(headers.get("Content-Length", 0))

            if content_length:
                message_text = STDIN.read(content_length)
                message_json = json.loads(message_text)
                our_user_input = message_json["result"]["userInput"]
                if message_json["id"] == request_id:
                    send_response(our_user_input, message_json["id"])
                elif message_json["method"] == "exit":
                    sys.exit(0)

        except Exception:
            print_log(traceback.format_exc())


def exec_function(user_input):
    try:
        compile(user_input, "", "eval")
    except SyntaxError:
        return exec
    return eval


def check_valid_command(request):
    try:
        user_input = request["params"]
        ast.parse(user_input[0])
        send_response("True", request["id"])
    except SyntaxError:
        send_response("False", request["id"])


def execute(request, user_globals):
    str_output = CustomIO("", encoding="utf-8")
    str_error = CustomIO("", encoding="utf-8")

    with redirect_io("stdout", str_output):
        with redirect_io("stderr", str_error):
            str_input = CustomIO("", encoding="utf-8", newline="\n")
            with redirect_io("stdin", str_input):
                exec_user_input(request["params"], user_globals)
    send_response(str_output.get_value(), request["id"])


def exec_user_input(user_input, user_globals):
    user_input = user_input[0] if isinstance(user_input, list) else user_input

    try:
        callable = exec_function(user_input)
        retval = callable(user_input, user_globals)
        if retval is not None:
            print(retval)
    except KeyboardInterrupt:
        print(traceback.format_exc())
    except Exception:
        print(traceback.format_exc())


class CustomIO(io.TextIOWrapper):
    """Custom stream object to replace stdio."""

    def __init__(self, name, encoding="utf-8", newline=None):
        self._buffer = io.BytesIO()
        self._custom_name = name
        super().__init__(self._buffer, encoding=encoding, newline=newline)

    def close(self):
        """Provide this close method which is used by some tools."""
        # This is intentionally empty.

    def get_value(self) -> str:
        """Returns value from the buffer as string."""
        self.seek(0)
        return self.read()


@contextlib.contextmanager
def redirect_io(stream: str, new_stream):
    """Redirect stdio streams to a custom stream."""
    old_stream = getattr(sys, stream)
    setattr(sys, stream, new_stream)
    yield
    setattr(sys, stream, old_stream)


def get_headers():
    headers = {}
    while line := STDIN.readline().strip():
        name, value = line.split(":", 1)
        headers[name] = value.strip()
    return headers


if __name__ == "__main__":
    while not STDIN.closed:
        try:
            headers = get_headers()
            content_length = int(headers.get("Content-Length", 0))

            if content_length:
                request_text = STDIN.read(content_length)
                request_json = json.loads(request_text)
                if request_json["method"] == "execute":
                    execute(request_json, USER_GLOBALS)
                if request_json["method"] == "check_valid_command":
                    check_valid_command(request_json)
                elif request_json["method"] == "exit":
                    sys.exit(0)

        except Exception:
            print_log(traceback.format_exc())