import pickle
import socket
import subprocess
import tempfile
import os
import zipfile
import io
import threading
import websockets
from .core import recv_msg, send_msg, TaskRequest, TaskResult, zip_dir


class BaseWorker:
    def __init__(self):
        self.base_command = 'python'


    def process_task(self, task: TaskRequest) -> TaskResult:
        with tempfile.TemporaryDirectory() as temp_dir:
            with zipfile.ZipFile(io.BytesIO(task.project_zip)) as zf:
                zf.extractall(temp_dir)

            script_dir = os.path.dirname(os.path.join(temp_dir, task.script_rel_path))
            script_name = os.path.basename(task.script_rel_path)

            cmd = task.run_cmd
            cmd[1] = script_name
            cmd[0] = self.base_command

            try:
                res_path = os.path.join(script_dir, task.result_dir)
                os.makedirs(res_path, exist_ok=True)

                proc = subprocess.run(
                    cmd,
                    cwd=script_dir,
                    capture_output=True,
                    text=True,
                    check=True
                )

                res_zip = zip_dir(res_path)

                return TaskResult(
                    task_id=task.task_id,
                    success=True,
                    stdout=proc.stdout,
                    stderr=proc.stderr,
                    result_zip=res_zip
                )
            except subprocess.CalledProcessError as e:
                return TaskResult(
                    task_id=task.task_id,
                    success=False,
                    stdout=e.stdout,
                    stderr=e.stderr
                )
            except Exception as e:
                return TaskResult(task_id=task.task_id, success=False, stdout="", stderr=str(e))


class SocketServer(BaseWorker):
    def __init__(self, port, password, command):
        super().__init__()
        self.port = port
        self.password = password
        self.base_command = command

    def start(self):
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.bind(('0.0.0.0', self.port))
        server.listen(5)
        print(f"Worker listening on {self.port}")

        while True:
            conn, addr = server.accept()
            t = threading.Thread(target=self.handle_client, args=(conn,))
            t.start()

    def handle_client(self, conn):
        try:
            auth = recv_msg(conn)
            if auth != self.password:
                print("Auth failed")
                conn.close()
                return

            send_msg(conn, "OK")

            task = recv_msg(conn)
            if isinstance(task, TaskRequest):
                result = self.process_task(task)
                send_msg(conn, result)
        except Exception as e:
            print(f"Error: {e}")
        finally:
            conn.close()


class CenterClientWorker(BaseWorker):
    def __init__(self, center_addr, password, base_command):
        self.uri = f"ws://{center_addr}/ws/register?password={password}"
        self.base_command = base_command

    async def run(self):
        async with websockets.connect(self.uri, ping_interval=None) as websocket:
            print(f"Connected to Center {self.uri}")
            while True:
                try:
                    message = await websocket.recv()
                    task = pickle.loads(message)
                    print(f"Processing task {task.task_id}...")

                    result = self.process_task(task)

                    await websocket.send(pickle.dumps(result))
                    print(f"Task {task.task_id} result sent.")

                except websockets.exceptions.ConnectionClosed:
                    print("Connection lost, retrying...")
                    break
                except Exception as e:
                    print(f"Worker Error: {e}")
                    err_res = TaskResult(task.task_id if 'task' in locals() else 'unknown', False, "", str(e))
                    await websocket.send(pickle.dumps(err_res))


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8080)
    parser.add_argument('--password', type=str, default='123456')
    parser.add_argument('--command', type=str, default='python')
    args = parser.parse_args()
    SocketServer(args.port, args.password, args.command).start()