import os
import socket
import threading
import base64
import random
import time


HOST = os.environ.get("CTF_HOST", "0.0.0.0")
PORT = int(os.environ.get("CTF_PORT", "9000"))
ROUNDS = int(os.environ.get("CTF_ROUNDS", "200"))
TIMEOUT_PER_ROUND = float(os.environ.get("CTF_TIMEOUT", "1.0"))  
WORKSPACE = os.getcwd()


def list_images(dir_path):
    if not os.path.isdir(dir_path):
        return []
    files = []
    for name in sorted(os.listdir(dir_path)):
        full = os.path.join(dir_path, name)
        if os.path.isfile(full):
            ext = os.path.splitext(full)[1].lower()
            if ext in {".jpg", ".jpeg", ".png", ".bmp"}:
                files.append(full)
    return files


def build_dataset(root):
    real_dir = os.path.join(root, "realface")
    fake_dir = os.path.join(root, "fakeface")
    reals = list_images(real_dir)
    fakes = list_images(fake_dir)
    if not reals or not fakes:
        raise RuntimeError("realface 或 fakeface 为空，请先准备数据")
    return reals, fakes


def send_line(conn, s: str):
    conn.sendall((s + "\n").encode())


def recv_line(conn, timeout: float | None = None):
    # timeout=None 表示不设置超时（阻塞等待）
    if timeout is not None:
        conn.settimeout(timeout)
    else:
        conn.settimeout(None)
    buf = b""
    while True:
        try:
            ch = conn.recv(1)
        except socket.timeout:
            raise
        if not ch:
            # 连接关闭
            return buf.decode(errors="ignore")
        if ch == b"\n":
            break
        buf += ch
    return buf.decode(errors="ignore")


def handle_client(conn, addr):
    try:
        reals, fakes = build_dataset(WORKSPACE)
        rng = random.Random()
        rng.seed(int(time.time()))
        send_line(conn, "欢迎来到存在伪人的世界！")
        send_line(conn, f"在这里，你将会遇到{ROUNDS}个人，你需要在 {TIMEOUT_PER_ROUND} 秒内判断他们是否是伪人。否则你将会被伪人同化！！！！！！！")
        send_line(conn, "你的正确率需要达到90%。")
        send_line(conn, "输入 confirm 开始挑战；每轮请在1秒内回复 0(真实)/1(虚假)")
        send_line(conn, "WAITING confirm...")

        while True:
            try:
                first = recv_line(conn, timeout=None).strip().lower()
            except socket.timeout:

                continue
            if not first:

                continue
            if first == "confirm":
                send_line(conn, "OK, start!")
                break
            else:
                send_line(conn, "请输入 confirm 开始挑战")

        correct = 0
        total = 0

        for i in range(1, ROUNDS + 1):
            lab = rng.choice([0, 1])
            src = reals if lab == 0 else fakes
            path = rng.choice(src)
            with open(path, "rb") as f:
                data = f.read()
            b64 = base64.b64encode(data).decode()

            send_line(conn, f"ROUND {i}")
            send_line(conn, "BEGIN")
            send_line(conn, b64)
            send_line(conn, "END")
            send_line(conn, "ANS?")


            try:
                ans = recv_line(conn, timeout=TIMEOUT_PER_ROUND).strip()
            except socket.timeout:
                send_line(conn, "你失败了")
                return

            if ans not in {"0", "1"}:
                send_line(conn, "你失败了")
                return

            total += 1
            if int(ans) == lab:
                correct += 1

        acc = (correct / total) if total > 0 else 0.0
        send_line(conn, f"结果: 正确={correct}, 总共={total}, 准确率={acc:.4f}")
        if acc >= 0.90:
            flag = os.environ.get("GZCTF_FLAG", "出错了")
            send_line(conn, f"恭喜你，你成功逃脱伪人的追捕 这是你的奖励：{flag}")
        else:
            send_line(conn, "Try again!")
    except Exception as e:
        try:
            send_line(conn, f"Server error: {e}")
        except Exception:
            pass
    finally:
        try:
            conn.shutdown(socket.SHUT_RDWR)
        except Exception:
            pass
        conn.close()


def main():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((HOST, PORT))
        s.listen(1)
        print(f"Challenge server listening on {HOST}:{PORT}")
        while True:
            conn, addr = s.accept()
            print(f"Client connected: {addr}")
            t = threading.Thread(target=handle_client, args=(conn, addr), daemon=True)
            t.start()


if __name__ == "__main__":
    main()