#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
7884 Fuzz Agent - 基于7884结构归纳引擎的模糊测试Agent
用于调用7884_brain.exe核心程序并提供Web界面
支持: 结构归纳分析、Gen样本生成(basic/pro)、Fuzz模糊测试、分析队列与记录管理
"""

import os
import sys
import json
import uuid
import time
import glob
import subprocess
import threading
import webbrowser
import tkinter as tk
from tkinter import filedialog
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from datetime import datetime

HOST = "127.0.0.1"
PORT = 7884
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
BRAIN_EXE = os.path.join(SCRIPT_DIR, "7884_brain.exe")
INDEX_HTML = os.path.join(SCRIPT_DIR, "index.html")
RECORDS_DIR = os.path.join(SCRIPT_DIR, "records")
FUZZ_DIR = os.path.join(SCRIPT_DIR, "fuzz_workspace")

os.makedirs(RECORDS_DIR, exist_ok=True)
os.makedirs(FUZZ_DIR, exist_ok=True)

task_queue = []
task_lock = threading.Lock()
task_history = {}
running_processes = {}
process_lock = threading.Lock()

fuzz_state = {
    "running": False,
    "target_exe": "",
    "sample_dir": "",
    "sample_list": [],
    "current_index": 0,
    "crashes": [],
    "total_run": 0,
    "process": None,
    "stop_flag": False,
    "log": []
}
fuzz_lock = threading.Lock()


def generate_task_id():
    return datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:6]


def save_record(task_id, record_data):
    path = os.path.join(RECORDS_DIR, f"{task_id}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(record_data, f, ensure_ascii=False, indent=2)


def load_records():
    records = []
    for fp in glob.glob(os.path.join(RECORDS_DIR, "*.json")):
        try:
            with open(fp, "r", encoding="utf-8") as f:
                records.append(json.load(f))
        except Exception:
            pass
    records.sort(key=lambda x: x.get("created_at", ""), reverse=True)
    return records


def delete_record(task_id):
    path = os.path.join(RECORDS_DIR, f"{task_id}.json")
    if os.path.exists(path):
        os.remove(path)
        return True
    return False


def parse_xml_default_output(xml_file):
    try:
        import xml.etree.ElementTree as ET
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for pub in root.iter('Publisher'):
            for p in pub.findall('Param'):
                if p.get('name') == 'FileName':
                    return p.get('value', '')
    except Exception:
        pass
    return ""


def find_peach_xml(output_dir):
    for root, dirs, files in os.walk(output_dir):
        for fn in files:
            if fn in ("format_schema_peach.xml", "format_schema.xml"):
                return os.path.join(root, fn).replace("\\", "/")
    return None


def worker_loop():
    while True:
        task = None
        with task_lock:
            for t in task_queue:
                if t["status"] == "pending":
                    t["status"] = "running"
                    t["started_at"] = datetime.now().isoformat()
                    task = t
                    break
        if not task:
            time.sleep(0.5)
            continue

        try:
            _execute_task(task)
        except Exception as e:
            task["status"] = "failed"
            task["error"] = str(e)
            task["finished_at"] = datetime.now().isoformat()
        save_record(task["task_id"], task)
        with task_lock:
            if task in task_queue:
                task_queue.remove(task)


def _execute_task(task):
    mode = task.get("mode", "analyze")

    if mode == "analyze":
        _run_analyze(task)
    elif mode == "gen":
        _run_gen(task)
    else:
        task["status"] = "failed"
        task["error"] = f"Unknown mode: {mode}"


def _run_analyze(task):
    if not os.path.exists(BRAIN_EXE):
        task["status"] = "failed"
        task["error"] = f"7884_brain.exe not found at: {BRAIN_EXE}"
        print(f"[ERROR] Brain exe not found: {BRAIN_EXE}")
        return

    input_dir = task.get("input", "")
    output_dir = task.get("output_dir", "")
    if not input_dir or not output_dir:
        task["status"] = "failed"
        task["error"] = "Missing input or output path"
        print(f"[ERROR] Missing path: input={input_dir}, output={output_dir}")
        return

    print(f"[Analyze] Input dir exists: {os.path.exists(input_dir)} -> {input_dir}")
    print(f"[Analyze] Output dir exists: {os.path.exists(output_dir)} -> {output_dir}")

    if not os.path.exists(input_dir):
        os.makedirs(input_dir, exist_ok=True)
        print(f"[Analyze] Created input dir: {input_dir}")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)
        print(f"[Analyze] Created output dir: {output_dir}")

    input_files = os.listdir(input_dir) if os.path.isdir(input_dir) else []
    print(f"[Analyze] Files in input dir: {len(input_files)} -> {input_files[:10]}")

    cmd = [BRAIN_EXE, "--input", input_dir, "--output", output_dir]

    run_mode = task.get("run_mode", "full")
    if run_mode == "cpu-only":
        cmd.extend(["--mode", "cpu-only"])
    elif run_mode == "lightweight":
        cmd.extend(["--mode", "lightweight"])

    if task.get("use_iffa", False):
        cmd.append("--use-iffa")
        iffa_path = task.get("iffa_path", "")
        if iffa_path:
            cmd.extend(["--iffa-path", iffa_path])

    max_iter = task.get("max_iterations", 0)
    if max_iter > 0:
        cmd.extend(["--max-iterations", str(max_iter)])

    patience = task.get("patience", 0)
    if patience > 0:
        cmd.extend(["--patience", str(patience)])

    if task.get("cuda_disable", False):
        cmd.append("--cuda-disable")

    cpu_threads = task.get("cpu_threads", 0)
    if cpu_threads > 0:
        cmd.extend(["--cpu-threads", str(cpu_threads)])

    cmd_str = ' '.join(f'"{c}"' if ' ' in c else c for c in cmd)
    print(f"[Analyze] === EXECUTING ===")
    print(f"[Analyze] Command: {cmd_str}")
    print(f"[Analyze] CWD: {os.getcwd()}")
    print(f"[Analyze] Brain size: {os.path.getsize(BRAIN_EXE)} bytes")

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
    )

    with process_lock:
        running_processes[task["task_id"]] = proc

    stdout_lines = []
    try:
        for line in proc.stdout:
            stdout_lines.append(line)
            print(f"[Brain] {line.rstrip()}")
    except Exception as e:
        print(f"[Analyze] Error reading output: {e}")

    proc.wait(timeout=7200)
    stdout = ''.join(stdout_lines)

    with process_lock:
        running_processes.pop(task["task_id"], None)

    task["returncode"] = proc.returncode
    task["stdout"] = stdout or ""

    print(f"[Analyze] Process exited with code: {proc.returncode}")
    if proc.returncode != 0:
        print(f"[Analyze] Last 20 lines of output:")
        for l in stdout_lines[-20:]:
            print(f"  | {l.rstrip()}")

    if proc.returncode == 0:
        task["status"] = "completed"
        task["peach_xml"] = find_peach_xml(output_dir)
        print(f"[Analyze] SUCCESS! Peach XML: {task['peach_xml']}")
    else:
        task["status"] = "failed"
        task["error"] = f"Process exited with code {proc.returncode}"

    task["finished_at"] = datetime.now().isoformat()


def _run_gen(task):
    if not os.path.exists(BRAIN_EXE):
        task["status"] = "failed"
        task["error"] = f"7884_brain.exe not found at: {BRAIN_EXE}"
        print(f"[ERROR] Brain exe not found: {BRAIN_EXE}")
        return

    xml_file = task.get("xml_file", "")
    if not xml_file or not os.path.exists(xml_file):
        task["status"] = "failed"
        task["error"] = f"Peach XML not found: {xml_file}"
        print(f"[ERROR] Peach XML not found: {xml_file}, exists={os.path.exists(xml_file) if xml_file else 'N/A'}")
        return

    cmd = [BRAIN_EXE, "--Gen", xml_file]

    iterations = task.get("iterations", 1000)
    cmd.extend(["-i", str(iterations)])

    output_template = task.get("output_template", "")
    if output_template:
        cmd.extend(["-o", output_template])

    seed = task.get("seed", 0)
    if seed > 0:
        cmd.extend(["-s", str(seed)])

    engine = task.get("engine", "pro")
    if engine == "pro":
        cmd.extend(["--engine", "pro"])

    range_val = task.get("range", "")
    if range_val:
        cmd.extend(["--range", range_val])

    cmd_str = ' '.join(f'"{c}"' if ' ' in c else c for c in cmd)
    print(f"[Gen] === EXECUTING ===")
    print(f"[Gen] Command: {cmd_str}")
    print(f"[Gen] XML exists: {os.path.exists(xml_file)}, size: {os.path.getsize(xml_file) if os.path.exists(xml_file) else 0}")

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
    )

    with process_lock:
        running_processes[task["task_id"]] = proc

    stdout_lines = []
    try:
        for line in proc.stdout:
            stdout_lines.append(line)
            print(f"[Gen-Brain] {line.rstrip()}")
    except Exception as e:
        print(f"[Gen] Error reading output: {e}")

    proc.wait(timeout=7200)
    stdout = ''.join(stdout_lines)

    with process_lock:
        running_processes.pop(task["task_id"], None)

    task["returncode"] = proc.returncode
    task["stdout"] = stdout or ""

    print(f"[Gen] Process exited with code: {proc.returncode}")
    if proc.returncode != 0:
        print(f"[Gen] Last 20 lines of output:")
        for l in stdout_lines[-20:]:
            print(f"  | {l.rstrip()}")

    task["status"] = "completed" if proc.returncode == 0 else "failed"
    if proc.returncode != 0:
        task["error"] = f"Gen process exited with code {proc.returncode}"
    else:
        print(f"[Gen] SUCCESS!")
    task["finished_at"] = datetime.now().isoformat()


def fuzz_worker():
    global fuzz_state
    while True:
        with fuzz_lock:
            if not fuzz_state["running"] or fuzz_state["stop_flag"]:
                fuzz_state["running"] = False
                fuzz_state["stop_flag"] = False
                return
            if fuzz_state["current_index"] >= len(fuzz_state["sample_list"]):
                fuzz_state["running"] = False
                fuzz_state["log"].append(f"[{datetime.now().isoformat()}] All samples processed, fuzzing complete")
                return

            sample_path = fuzz_state["sample_list"][fuzz_state["current_index"]]
            fuzz_state["current_index"] += 1
            target_exe = fuzz_state["target_exe"]

        fuzz_state["log"].append(f"[{datetime.now().isoformat()}] Loading sample {fuzz_state['current_index']}/{len(fuzz_state['sample_list'])}: {os.path.basename(sample_path)}")

        try:
            start_time = time.time()
            proc = subprocess.Popen(
                [target_exe, sample_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
            )
            fuzz_state["process"] = proc

            try:
                return_code = proc.wait(timeout=30)
                elapsed = time.time() - start_time
            except subprocess.TimeoutExpired:
                proc.kill()
                proc.wait()
                return_code = -1
                elapsed = 30.0

            fuzz_state["total_run"] += 1

            if return_code != 0 and return_code != 1:
                crash_info = {
                    "sample": sample_path,
                    "sample_name": os.path.basename(sample_path),
                    "return_code": return_code,
                    "elapsed": round(elapsed, 2),
                    "timestamp": datetime.now().isoformat(),
                    "index": fuzz_state["total_run"]
                }
                fuzz_state["crashes"].append(crash_info)
                fuzz_state["log"].append(f"[{datetime.now().isoformat()}] CRASH detected! Sample: {os.path.basename(sample_path)}, Exit code: {return_code}")
            else:
                fuzz_state["log"].append(f"[{datetime.now().isoformat()}] Sample OK (exit={return_code}, {elapsed:.1f}s)")

        except Exception as e:
            fuzz_state["log"].append(f"[{datetime.now().isoformat()}] Error running sample: {e}")
            fuzz_state["total_run"] += 1

        time.sleep(0.1)


class BrainHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        print(f"[{self.address_string()}] {format % args}")

    def _send_json(self, data, status=200):
        content = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", len(content))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(content)

    def _send_file(self, filepath, content_type="text/html"):
        if not os.path.exists(filepath):
            self.send_error(404, "File not found")
            return
        with open(filepath, "rb") as f:
            content = f.read()
        self.send_response(200)
        self.send_header("Content-Type", f"{content_type}; charset=utf-8")
        self.send_header("Content-Length", len(content))
        self.end_headers()
        self.wfile.write(content)

    def _read_body(self):
        content_length = int(self.headers.get("Content-Length", 0))
        return self.rfile.read(content_length).decode("utf-8")

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path

        if path == "/" or path == "/index.html":
            self._send_file(INDEX_HTML)
        elif path == "/api/status":
            self._handle_status()
        elif path == "/api/select-folder":
            self._select_folder()
        elif path == "/api/select-file":
            self._select_file()
        elif path == "/api/queue":
            self._handle_queue()
        elif path == "/api/records":
            self._handle_records()
        elif path.startswith("/api/records/"):
            task_id = path.split("/api/records/")[1]
            self._handle_record_detail(task_id)
        elif path == "/api/fuzz/status":
            self._handle_fuzz_status()
        elif path.startswith("/api/output/"):
            rel_path = path[len("/api/output/"):]
            self._serve_output_file(rel_path)
        else:
            self.send_error(404, "Not found")

    def do_POST(self):
        parsed = urlparse(self.path)
        path = parsed.path

        if path == "/api/run":
            self._handle_run()
        elif path == "/api/gen":
            self._handle_gen()
        elif path == "/api/queue/cancel":
            self._handle_queue_cancel()
        elif path == "/api/fuzz/start":
            self._handle_fuzz_start()
        elif path == "/api/fuzz/stop":
            self._handle_fuzz_stop()
        elif path == "/api/fuzz/generate-samples":
            self._handle_fuzz_generate_samples()
        else:
            self.send_error(404, "Not found")

    def do_DELETE(self):
        parsed = urlparse(self.path)
        path = parsed.path

        if path.startswith("/api/records/"):
            task_id = path.split("/api/records/")[1]
            ok = delete_record(task_id)
            self._send_json({"success": ok})
        else:
            self.send_error(404, "Not found")

    def _handle_status(self):
        with task_lock:
            pending = sum(1 for t in task_queue if t["status"] == "pending")
            running = sum(1 for t in task_queue if t["status"] == "running")

        with fuzz_lock:
            fuzz_running = fuzz_state["running"]

        self._send_json({
            "status": "ok",
            "brain_exists": os.path.exists(BRAIN_EXE),
            "queue": {"pending": pending, "running": running, "total": len(task_queue)},
            "fuzz_running": fuzz_running
        })

    def _select_folder(self):
        root = tk.Tk()
        root.withdraw()
        root.attributes("-topmost", True)
        folder = filedialog.askdirectory(title="选择目录")
        root.destroy()
        self._send_json({"path": folder if folder else ""})

    def _select_file(self):
        parsed = urlparse(self.path)
        qs = parse_qs(parsed.query)
        file_type = qs.get("type", ["exe"])[0]

        if file_type == "xml":
            filetypes = [
                ("XML文件", "*.xml"),
                ("所有文件", "*.*")
            ]
            title = "选择XML文件"
        elif file_type == "exe":
            filetypes = [
                ("可执行文件", "*.exe"),
                ("所有文件", "*.*")
            ]
            title = "选择可执行文件"
        else:
            filetypes = [
                ("所有文件", "*.*")
            ]
            title = "选择文件"

        root = tk.Tk()
        root.withdraw()
        root.attributes("-topmost", True)
        file_path = filedialog.askopenfilename(title=title, filetypes=filetypes)
        root.destroy()
        self._send_json({"path": file_path if file_path else ""})

    def _handle_queue(self):
        with task_lock:
            items = []
            for t in task_queue:
                items.append({
                    "task_id": t["task_id"],
                    "mode": t.get("mode", ""),
                    "status": t["status"],
                    "created_at": t.get("created_at", ""),
                    "started_at": t.get("started_at", ""),
                    "finished_at": t.get("finished_at", ""),
                    "error": t.get("error", ""),
                    "label": t.get("label", "")
                })
        self._send_json({"queue": items})

    def _handle_records(self):
        records = load_records()
        self._send_json({"records": records})

    def _handle_record_detail(self, task_id):
        path = os.path.join(RECORDS_DIR, f"{task_id}.json")
        if not os.path.exists(path):
            self._send_json({"error": "Record not found"}, 404)
            return
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        self._send_json(data)

    def _handle_run(self):
        body = self._read_body()
        try:
            params = json.loads(body)
        except json.JSONDecodeError:
            self._send_json({"error": "Invalid JSON"}, 400)
            return

        task_id = generate_task_id()
        task = {
            "task_id": task_id,
            "mode": "analyze",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "started_at": "",
            "finished_at": "",
            "input": params.get("input", ""),
            "output_dir": params.get("output", ""),
            "run_mode": params.get("run_mode", "full"),
            "use_iffa": params.get("use_iffa", False),
            "iffa_path": params.get("iffa_path", ""),
            "max_iterations": params.get("max_iterations", 0),
            "patience": params.get("patience", 0),
            "cuda_disable": params.get("cuda_disable", False),
            "cpu_threads": params.get("cpu_threads", 0),
            "label": params.get("label", ""),
            "stdout": "",
            "error": "",
            "returncode": None,
            "peach_xml": ""
        }

        with task_lock:
            task_queue.append(task)
        save_record(task_id, task)

        print(f"[API] New analyze task: {task_id}")
        print(f"[API]   input={params.get('input', '')}")
        print(f"[API]   output={params.get('output', '')}")
        print(f"[API]   run_mode={params.get('run_mode', 'full')}")
        print(f"[API]   label={params.get('label', '')}")

        self._send_json({"success": True, "task_id": task_id})

    def _handle_gen(self):
        body = self._read_body()
        try:
            params = json.loads(body)
        except json.JSONDecodeError:
            self._send_json({"error": "Invalid JSON"}, 400)
            return

        xml_file = params.get("xml_file", "")
        if not xml_file:
            self._send_json({"error": "xml_file is required"}, 400)
            return

        task_id = generate_task_id()
        output_template = params.get("output_template", "")
        gen_output_dir = ""
        if not output_template:
            xml_default = parse_xml_default_output(xml_file)
            if xml_default:
                gen_output_dir = os.path.dirname(xml_default.replace("{0}", "0"))
                os.makedirs(gen_output_dir, exist_ok=True)
            output_template = ""

        task = {
            "task_id": task_id,
            "mode": "gen",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "started_at": "",
            "finished_at": "",
            "xml_file": xml_file,
            "iterations": params.get("iterations", 1000),
            "output_template": output_template,
            "seed": params.get("seed", 0),
            "engine": params.get("engine", "pro"),
            "range": params.get("range", ""),
            "label": params.get("label", ""),
            "output": "",
            "error": "",
            "returncode": None
        }

        with task_lock:
            task_queue.append(task)
        save_record(task_id, task)

        self._send_json({"success": True, "task_id": task_id, "output_dir": gen_output_dir})

    def _handle_queue_cancel(self):
        body = self._read_body()
        try:
            params = json.loads(body)
        except json.JSONDecodeError:
            self._send_json({"error": "Invalid JSON"}, 400)
            return

        task_id = params.get("task_id", "")
        with task_lock:
            for t in task_queue:
                if t["task_id"] == task_id:
                    if t["status"] == "pending":
                        t["status"] = "cancelled"
                        t["finished_at"] = datetime.now().isoformat()
                        save_record(task_id, t)
                        self._send_json({"success": True, "message": "Task cancelled"})
                        return
                    elif t["status"] == "running":
                        with process_lock:
                            proc = running_processes.get(task_id)
                            if proc:
                                try:
                                    proc.kill()
                                except Exception:
                                    pass
                        t["status"] = "cancelled"
                        t["finished_at"] = datetime.now().isoformat()
                        save_record(task_id, t)
                        self._send_json({"success": True, "message": "Task killed"})
                        return

        self._send_json({"success": False, "error": "Task not found or not cancellable"})

    def _handle_fuzz_start(self):
        global fuzz_state
        body = self._read_body()
        try:
            params = json.loads(body)
        except json.JSONDecodeError:
            self._send_json({"error": "Invalid JSON"}, 400)
            return

        with fuzz_lock:
            if fuzz_state["running"]:
                self._send_json({"error": "Fuzz already running"}, 400)
                return

            target_exe = params.get("target_exe", "")
            sample_dir = params.get("sample_dir", "")

            if not target_exe or not os.path.exists(target_exe):
                self._send_json({"error": "Target executable not found"}, 400)
                return
            if not sample_dir or not os.path.isdir(sample_dir):
                self._send_json({"error": "Sample directory not found"}, 400)
                return

            sample_list = []
            for ext in ["*.*"]:
                for fp in glob.glob(os.path.join(sample_dir, ext)):
                    if os.path.isfile(fp):
                        sample_list.append(fp)
            sample_list.sort()

            if not sample_list:
                self._send_json({"error": "No samples found in directory"}, 400)
                return

            fuzz_state = {
                "running": True,
                "target_exe": target_exe,
                "sample_dir": sample_dir,
                "sample_list": sample_list,
                "current_index": 0,
                "crashes": [],
                "total_run": 0,
                "process": None,
                "stop_flag": False,
                "log": [f"[{datetime.now().isoformat()}] Fuzzing started: {len(sample_list)} samples, target={os.path.basename(target_exe)}"]
            }

        t = threading.Thread(target=fuzz_worker, daemon=True)
        t.start()

        self._send_json({"success": True, "message": f"Fuzzing started with {len(sample_list)} samples"})

    def _handle_fuzz_stop(self):
        global fuzz_state
        with fuzz_lock:
            if not fuzz_state["running"]:
                self._send_json({"error": "Fuzz not running"}, 400)
                return
            fuzz_state["stop_flag"] = True
            if fuzz_state["process"]:
                try:
                    fuzz_state["process"].kill()
                except Exception:
                    pass
            fuzz_state["log"].append(f"[{datetime.now().isoformat()}] Fuzzing stopped by user")

        self._send_json({"success": True, "message": "Fuzz stopping..."})

    def _handle_fuzz_status(self):
        with fuzz_lock:
            state_copy = {
                "running": fuzz_state["running"],
                "target_exe": fuzz_state["target_exe"],
                "sample_dir": fuzz_state["sample_dir"],
                "total_samples": len(fuzz_state["sample_list"]),
                "current_index": fuzz_state["current_index"],
                "crashes": list(fuzz_state["crashes"]),
                "total_run": fuzz_state["total_run"],
                "log": fuzz_state["log"][-50:]
            }
        self._send_json(state_copy)

    def _handle_fuzz_generate_samples(self):
        body = self._read_body()
        try:
            params = json.loads(body)
        except json.JSONDecodeError:
            self._send_json({"error": "Invalid JSON"}, 400)
            return

        xml_file = params.get("xml_file", "")
        if not xml_file or not os.path.exists(xml_file):
            self._send_json({"error": "Peach XML not found"}, 400)
            return

        iterations = params.get("iterations", 1000)
        engine = params.get("engine", "pro")
        seed = params.get("seed", 0)

        task_id = generate_task_id()
        gen_output_dir = ""
        xml_default = parse_xml_default_output(xml_file)
        if xml_default:
            gen_output_dir = os.path.dirname(xml_default.replace("{0}", "0"))
            os.makedirs(gen_output_dir, exist_ok=True)
        output_template = ""

        task = {
            "task_id": task_id,
            "mode": "gen",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "started_at": "",
            "finished_at": "",
            "xml_file": xml_file,
            "iterations": iterations,
            "output_template": output_template,
            "seed": seed,
            "engine": engine,
            "range": "",
            "label": f"Fuzz样本生成-{engine}",
            "output": "",
            "error": "",
            "returncode": None
        }

        with task_lock:
            task_queue.append(task)
        save_record(task_id, task)

        self._send_json({
            "success": True,
            "task_id": task_id,
            "output_dir": gen_output_dir,
            "message": f"Sample generation queued: {iterations} iterations, engine={engine}"
        })

    def _serve_output_file(self, rel_path):
        safe_path = os.path.normpath(rel_path)
        if safe_path.startswith("..") or ":" in safe_path:
            self.send_error(403, "Forbidden")
            return

        abs_path = os.path.join(SCRIPT_DIR, safe_path)
        if not os.path.exists(abs_path) or not os.path.isfile(abs_path):
            self.send_error(404, "File not found")
            return

        ext = os.path.splitext(abs_path)[1].lower()
        content_types = {
            ".html": "text/html",
            ".json": "application/json",
            ".xml": "application/xml",
            ".txt": "text/plain",
            ".py": "text/plain",
            ".csv": "text/csv",
            ".log": "text/plain"
        }
        ct = content_types.get(ext, "application/octet-stream")
        self._send_file(abs_path, ct)


def cleanup_processes():
    with process_lock:
        for proc in running_processes.values():
            try:
                proc.kill()
            except Exception:
                pass
        running_processes.clear()

    with task_lock:
        for t in task_queue:
            if t["status"] in ("pending", "running"):
                t["status"] = "cancelled"
                t["finished_at"] = datetime.now().isoformat()
                save_record(t["task_id"], t)

    with fuzz_lock:
        if fuzz_state["running"]:
            fuzz_state["stop_flag"] = True
            if fuzz_state["process"]:
                try:
                    fuzz_state["process"].kill()
                except Exception:
                    pass


def run_server():
    worker_thread = threading.Thread(target=worker_loop, daemon=True)
    worker_thread.start()

    server = HTTPServer((HOST, PORT), BrainHandler)
    print(f"7884 Fuzz Agent running at http://{HOST}:{PORT}/")
    print(f"Brain: {BRAIN_EXE}")
    print(f"Records: {RECORDS_DIR}")
    print("Press Ctrl+C to stop...")

    def open_browser():
        webbrowser.open(f"http://{HOST}:{PORT}/")

    threading.Timer(1.0, open_browser).start()

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")
    finally:
        cleanup_processes()
        server.shutdown()


if __name__ == "__main__":
    run_server()
