import os
import zipfile
import pickle
import struct
import io
from dataclasses import dataclass
from typing import Dict, Any, List

HEADER_SIZE = 4

def send_msg(sock, msg):
    msg_data = pickle.dumps(msg)
    sock.sendall(struct.pack('>I', len(msg_data)) + msg_data)

def recv_msg(sock):
    raw_msglen = recvall(sock, HEADER_SIZE)
    if not raw_msglen: return None
    msglen = struct.unpack('>I', raw_msglen)[0]
    data = recvall(sock, msglen)
    return pickle.loads(data)

def recvall(sock, n):
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet: return None
        data += packet
    return data

def zip_dir(path):
    mem_zip = io.BytesIO()
    with zipfile.ZipFile(mem_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
        for root, dirs, files in os.walk(path):
            for file in files:
                abs_path = os.path.join(root, file)
                rel_path = os.path.relpath(abs_path, os.path.join(path, '..'))
                zf.write(abs_path, rel_path)
    return mem_zip.getvalue()

@dataclass
class TaskRequest:
    task_id: str
    project_zip: bytes
    script_rel_path: str
    run_cmd: List[str]
    result_dir: str

@dataclass
class TaskResult:
    task_id: str
    success: bool
    stdout: str
    stderr: str
    result_zip: bytes = None