client.py

import json, re, socket, sys, time, requests, select, threading, queue
import multiprocessing as mp
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal, Optional

def encode_message(message: dict[str, Any]) -> bytes:
    return json.dumps(message).encode("utf-8")

def decode_message(data: bytes) -> dict[str, Any]:
    return json.loads(data.decode("utf-8"))

def send_message(connection: socket.socket, data: dict[str, Any]):
    """Send a JSON message to the server. For robustness across servers,
    append a single newline so implementations that use line-based reads don't block.
    Servers that parse raw JSON will ignore trailing whitespace/newline.
    """
    msg = encode_message(data) + b"\n"
    connection.sendall(msg)

def receive_message(connection: socket.socket, client: "Client") -> Optional[dict[str, Any]]:
    """Read one newline-delimited server message from a persistent buffer.
    Returns None if the socket is closed or decoding fails."""
    while True:
        nl = client._buf.find(b"\n")
        if nl != -1:
            line = client._buf[:nl]
            del client._buf[:nl+1]
            try:
                return decode_message(bytes(line))
            except (json.JSONDecodeError, UnicodeDecodeError):
                return None
        try:
            chunk: bytes = connection.recv(4096)
        except (ConnectionResetError, OSError):
            return None
        if not chunk:
            return None
        client._buf.extend(chunk)

def connect(username:str, host:str, port:int) -> Optional[socket.socket]:
    # AF_INET: IPv4 ; SOCK_STREAM: TCP
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect((host, port))
    except (OSError, socket.gaierror):
        # Per spec: print exactly "Connection failed" and then wait for proper input again
        print("Connection failed")
        try:
            s.close()
        except Exception:
            pass
        return None


    hi_dict = {
        "message_type": "HI",
        "username": username,
    }
    # Send HI (with newline terminator for compatibility)
    send_message(s, hi_dict)
    return s

def disconnect(connection: Optional[socket.socket]):
    if not connection:
        return
    try:
        send_message(connection, {"message_type":"BYE"})
    except (ConnectionError, OSError):
        pass
    try:
        connection.close()
    except (ConnectionError, OSError):
        pass


def answer_question(
    question_type: str,
    question: str,
    short_question: str,
    client_mode: Literal["you", "auto", "ai"],
    time_limit: float,
    client: "Client",
) -> Optional[str]:
    """Return an answer string or None on timeout/skip."""

    if client_mode == "you":
        # On POSIX, prefer select on stdin to avoid background threads that may
        # consume later commands (e.g., EXIT) after a timeout.
        if sys.platform != "win32":
            try:
                rlist, _, _ = select.select([sys.stdin], [], [], time_limit)
            except (ValueError, OSError):
                rlist = []
            if not rlist:
                return None
            try:
                line = sys.stdin.readline()
            except Exception:
                line = ""
            ans = line.strip()
            if ans.upper() == "DISCONNECT":
                client.request_disconnect = True
                return None
            if ans.upper() == "EXIT":
                client.request_exit = True
                return None
            return ans
        else:
            # Windows fallback: use a short-lived background thread with queue
            q: queue.Queue[str] = queue.Queue(maxsize=1)

            def _read_input():
                try:
                    text = input().strip()
                except EOFError:
                    text = ""
                q.put(text)

            t = threading.Thread(target=_read_input, daemon=True)
            t.start()
            try:
                ans = q.get(timeout=time_limit)
            except queue.Empty:
                return None
            if ans.upper() == "DISCONNECT":
                client.request_disconnect = True
                return None
            if ans.upper() == "EXIT":
                client.request_exit = True
                return None
            return ans
    elif client_mode == "auto":
        return auto_compute_answer_mp(question_type, short_question, time_limit)
    elif client_mode == "ai":
        return answer_question_ollama_mp(question, client.ollama_config, time_limit)
    return None

# -------------- Multiprocessing variants for concurrency --------------

def _auto_worker(question_type: str, short_question: str, out_q: "mp.Queue[str]"):
    try:
        if question_type == "Mathematics":
            val = eval_math_expression(short_question)
            
            if val is None:
                out_q.put(None)
            else:
                if float(val).is_integer():
                    out_q.put(str(int(round(val))))
                else:
                    out_q.put(str(val))
        elif question_type == "Roman Numerals":
            out_q.put(str(roman_to_int(short_question)))
        elif question_type == "Usable IP Addresses of a Subnet":
            out_q.put(str(usable_ips(short_question)))
        elif question_type == "Network and Broadcast Address of a Subnet":
            net, bcast = network_broadcast(short_question)
            out_q.put(f"{net} and {bcast}")
        else:
            out_q.put("")
    except (ValueError, KeyError, IndexError):
        out_q.put(None)


def auto_compute_answer_mp(question_type: str, short_question: str, time_limit: float) -> Optional[str]:
    out_q: mp.Queue = mp.Queue(maxsize=1)
    p = mp.Process(target=_auto_worker, args=(question_type, short_question, out_q), daemon=True)
    p.start()
    try:
        ans = out_q.get(timeout=time_limit)
    except queue.Empty:
        ans = None
    finally:
        if p.is_alive():
            p.terminate()
        p.join(timeout=0.2)
    return ans


def _ai_worker(question: str, ollama_config: dict, time_limit: float, out_q: "mp.Queue[str]"):
    try:
        host = ollama_config.get("ollama_host")
        port = ollama_config.get("ollama_port")
        model = ollama_config.get("ollama_model")
        url = f"http://{host}:{port}/api/chat"
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": question + " return just the answer part only. don't respond as a whole paragraph."}],
            "stream": False,
        }
        resp = requests.post(url, json=payload, timeout=time_limit)
        j = resp.json()
        msg = j.get("message") or {}
        content = msg.get("content")
        if isinstance(content, str):
            out_q.put(content)
            return
        if isinstance(j.get("messages"), list) and j["messages"]:
            m0 = j["messages"][0]
            if isinstance(m0, dict) and isinstance(m0.get("content"), str):
                out_q.put(m0["content"]) 
                return
        out_q.put(None)
    except (requests.RequestException, ValueError, KeyError):
        out_q.put(None)


def answer_question_ollama_mp(question: str, ollama_config: Optional[dict], time_limit: float) -> Optional[str]:
    if not ollama_config:
        return None
    out_q: mp.Queue = mp.Queue(maxsize=1)
    p = mp.Process(target=_ai_worker, args=(question, ollama_config, time_limit, out_q), daemon=True)
    p.start()
    try:
        ans = out_q.get(timeout=time_limit)
    except queue.Empty:
        ans = None
    finally:
        if p.is_alive():
            p.terminate()
        p.join(timeout=0.2)
    return ans

def eval_math_expression(expr: str) -> Optional[float]:
    try:
        tokens = tokenize_expr(expr)
        rpn = infix_to_rpn(tokens)
        return eval_rpn(rpn)
    except (ValueError, IndexError):
        return None

def tokenize_expr(expr: str) -> list[str]:
    tokens: list[str] = []
    i = 0
    while i < len(expr):
        c = expr[i]
        if c.isspace():
            i += 1
            continue
        if c.isdigit():
            j = i
            while j < len(expr) and expr[j].isdigit():
                j += 1
            tokens.append(expr[i:j])
            i = j
            continue
        if c in "+-*/()":
            tokens.append(c)
            i += 1
            continue
        raise ValueError("Invalid character in expression")
    return tokens

def infix_to_rpn(tokens: list[str]) -> list[str]:
    prec = {"+": 1, "-": 1, "*": 2, "/": 2}
    output: list[str] = []
    op_stack: list[str] = []
    for t in tokens:
        if t.isdigit():
            output.append(t)
        elif t in "+-*/":
            while op_stack and op_stack[-1] in prec and prec[op_stack[-1]] >= prec[t]:
                output.append(op_stack.pop())
            op_stack.append(t)
        elif t == "(":
            op_stack.append(t)
        elif t == ")":
            while op_stack and op_stack[-1] != "(":
                output.append(op_stack.pop())
            if not op_stack or op_stack[-1] != "(":
                raise ValueError("Mismatched parentheses")
            op_stack.pop()
        else:
            raise ValueError("Invalid token")
    while op_stack:
        op = op_stack.pop()
        if op in ("(", ")"):
            raise ValueError("Mismatched parentheses at end")
        output.append(op)
    return output

def eval_rpn(tokens: list[str]) -> float:
    stack: list[float] = []
    for t in tokens:
        if t.isdigit():
            stack.append(float(t))
        elif t in "+-*/":
            if len(stack) < 2:
                raise ValueError("Insufficient values")
            b = stack.pop()
            a = stack.pop()
            if t == "+":
                stack.append(a + b)
            elif t == "-":
                stack.append(a - b)
            elif t == "*":
                stack.append(a * b)
            elif t == "/":
                stack.append(a / b)
        else:
            raise ValueError("Invalid RPN token")
    if len(stack) != 1:
        raise ValueError("Invalid RPN evaluation")
    return stack[0]

def roman_to_int(s: str) -> int:
    values = {"I":1, "V":5, "X":10, "L":50, "C":100, "D":500, "M":1000}
    total = 0
    i = 0
    s = s.strip().upper()
    while i < len(s):
        v = values.get(s[i], 0)
        if i + 1 < len(s):
            v2 = values.get(s[i+1], 0)
            if v < v2:
                total += (v2 - v)
                i += 2
                continue
        total += v
        i += 1
    return total

def ip_to_int(ip: str) -> int:
    parts = [int(p) for p in ip.split(".")]
    n = 0
    for p in parts:
        n = (n << 8) | (p & 0xFF)
    return n

def int_to_ip(n: int) -> str:
    return ".".join(str((n >> (24 - 8*i)) & 0xFF) for i in range(4))

def usable_ips(cidr: str) -> int:
    ip, prefix_str = cidr.split("/")
    prefix = int(prefix_str)
    host_bits = 32 - prefix
    if prefix >= 32:
        return 1 if prefix == 32 else 0
    if prefix == 31:
        return 0
    total = (1 << host_bits)
    return max(0, total - 2)

def network_broadcast(cidr: str) -> tuple[str, str]:
    ip, prefix_str = cidr.split("/")
    prefix = int(prefix_str)
    ipn = ip_to_int(ip)
    mask = (0xFFFFFFFF << (32 - prefix)) & 0xFFFFFFFF if prefix > 0 else 0
    network = ipn & mask
    broadcast = network | (~mask & 0xFFFFFFFF)
    return int_to_ip(network), int_to_ip(broadcast)

def handle_command(username: str,
                   connection : Optional[socket.socket]) -> str | socket.socket:
    '''returns a command-type confirmation String code or a socket'''
    command:str = input()   # stdin input (no extra prompt text)
    if (m:= re.match(r"CONNECT (?P<hostname>[^:]+):(?P<port>\d+)", command)):
        # returns the socket object if connection successful
        return connect(username, m.group('hostname'), int(m.group('port')))
    if command.lower() == "disconnect":
        if connection:
            disconnect(connection)
            return "DISCONNECTED"
    if command.lower() == "exit": # program termination command
        if connection:
            disconnect(connection)
        sys.exit(0)
    return "" # empty string means handler should be delegated to answering questions

def handle_received_message(message: dict[str, Any],
                            client: "Client",
                            connection: socket.socket) -> bool:
    """Handles different kinds of messages sent by the Trivia.NET server.
    Returns True if game should continue, False if finished."""
    mtype = message.get("message_type")
    if mtype == "READY":
        # Client prints "info" on receive
        print(message.get("info", ""))
    elif mtype == "QUESTION":
        qtype = message.get("question_type")
        trivia_question = message.get("trivia_question")
        short_question = message.get("short_question")
        time_limit = message.get("time_limit")
        try:
            tl = float(time_limit)
        except (ValueError, TypeError):
            tl = 5.0
        # Client prints "trivia_question" upon receipt
        print(trivia_question)
        ans = answer_question(str(qtype), str(trivia_question), str(short_question), client.client_mode, tl, client)
        if client.request_exit:
            disconnect(connection)
            sys.exit(0)
        if client.request_disconnect:
            disconnect(connection)
            return False
        if ans is not None:
            out = {"message_type": "ANSWER", "answer": ans}
            send_message(connection, out)
    elif mtype == "RESULT":
        # Client prints the "feedback" value upon receipt
        print(message.get("feedback", ""))
        # No skip-next rule: always attempt every question
    elif mtype == "LEADERBOARD":
        # Client prints "state" when received
        print(message.get("state", ""))
    elif mtype == "FINISHED":
        # Client prints "final_standings" on receipt
        print(message.get("final_standings", ""))
        return False
    return True


@dataclass
class Client:
    '''
    Dataclass representing the client object. Attributes are the same as the config keys.
    '''
    username : str
    client_mode : Literal["you", "auto", "ai"]
    ollama_config : Optional[dict]
    request_disconnect: bool = False
    request_exit: bool = False
    _buf: bytearray = field(default_factory=bytearray)  # persistent newline framing buffer

def main():
    # Answer questions depending on the mode ('you', 'auto' or 'ai')

    # ===== arguments check =====
    argv = sys.argv[1:]
    if not argv or len(argv) < 2:
        print("client.py: Configuration not provided", file=sys.stderr)
        sys.exit(1)
    if argv[0] != "--config":
        print("client.py: Configuration not provided", file=sys.stderr)
        sys.exit(1)
    cfg_path = argv[1]
    p = Path(cfg_path)
    if not (p.exists() and p.is_file()):
        print(f"client.py: File {cfg_path} does not exist", file=sys.stderr)
        sys.exit(1)

    # ===== loading and parsing the config file =====
    with open(cfg_path,'r') as cfg_file:
        config = json.load(cfg_file)    # a dict is returned

    # ===== creating a Client object with the parsed config =====
    client = Client(username=config.get("username"),
                    client_mode=config.get("client_mode"),
                    ollama_config=config.get("ollama_config", None))


    # === client mode ai was chosen ===
    if client.client_mode == "ai":
        if client.ollama_config is None:
            print("client.py: Missing values for Ollama configuration", file=sys.stderr)
            sys.exit(1)


    # ===== Handle CONNECT, DISCONNECT, EXIT commands =====
    connection: Optional[socket.socket] = None
    while True:  # wait for CONNECT
        result = handle_command(username=client.username, connection=connection)

        if isinstance(result, socket.socket):
            connection = result
            # Reset any prior command flags when starting a new game/connection
            client.request_disconnect = False
            client.request_exit = False
        elif isinstance(result, str) and result == "DISCONNECTED":
            connection = None
            # Clear flags after a manual disconnect so the next game doesn't auto-disconnect
            client.request_disconnect = False
            client.request_exit = False
            continue
        else:
            # not connected yet, keep waiting for valid CONNECT
            continue

        # Connected: process messages until finished or disconnect
        while connection:
            msg = receive_message(connection, client)
            if msg is None:
                # Server stopped; disconnect self
                disconnect(connection)
                connection = None
                break
            if not handle_received_message(msg, client, connection):
                connection = None
                break


if __name__ == "__main__":
    main()