#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import subprocess
import time
from pathlib import Path


def read_text(path: Path, limit: int = 20000) -> str:
    if not path.exists():
        return ""
    text = path.read_text(encoding="utf-8", errors="replace")
    if len(text) <= limit:
        return text
    return text[-limit:]


def write_json(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def append_jsonl(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(value, ensure_ascii=False, sort_keys=True) + "\n")


def next_iteration(research: Path) -> int:
    highest = 0
    for folder in ("logs", "prompts", "runs"):
        path = research / folder
        if not path.exists():
            continue
        for item in path.iterdir():
            try:
                highest = max(highest, int(item.stem))
            except ValueError:
                pass
    return highest + 1


def build_prompt(workspace: Path, iteration: int) -> str:
    research = workspace / ".dumont_research"
    brief = read_text(workspace / "docs" / "autoresearch_tts_foundation_brief.md", 24000)
    qwen = read_text(workspace / "docs" / "qwen3_tts_smoke_test.md", 24000)
    memory = read_text(research / "wiki" / "memory.md", 24000)
    experiments = read_text(research / "experiments.jsonl", 16000)
    return f"""
You are running Dumont Rust AutoResearch iteration {iteration}.

Use model judgment, but act like a measured ML researcher.

Hard constraints:
- Use the Rust Dumont tool loop only.
- Do not add TypeScript backend code.
- Generated code must follow the project style: no comments, no docstrings, simplest correct change.
- Work inside {workspace}.
- Keep changes small and reversible.
- Prefer creating or improving benchmarks, manifests, prompts, evaluation scripts, and experiment plans before launching expensive training.
- Use web_search or WebFetch for at least one current external source unless the previous iteration already fetched the same source.
- Read local files before proposing changes.
- Record the result in .dumont_research/runs/{iteration:04d}.md.
- Append a compact JSON object to .dumont_research/experiments.jsonl.
- Update .dumont_research/wiki/memory.md with durable learnings.

Research question:
How can Qwen/Qwen3.5-0.8B drive Qwen/Qwen3-TTS-Tokenizer-12Hz tokens for faster, clearer Brazilian Portuguese teacher audio, using the validated project facts, fine-tuning fundamentals, and current literature?

Focus this iteration on exactly one useful step:
1. strengthen the benchmark for the student/teacher audio loop,
2. improve the hypothesis space for TTS token generation,
3. prepare a safe small experiment,
4. implement a small evaluation helper,
5. or run a cheap smoke/eval if dependencies are available.

Foundation brief:
{brief}

Qwen smoke doc:
{qwen}

Current memory:
{memory}

Recent experiments:
{experiments}

End with a short status block containing:
iteration, files_changed, command_run, metric, keep, next_step.
""".strip()


def run_iteration(args: argparse.Namespace, iteration: int) -> int:
    workspace = Path(args.workspace).resolve()
    research = workspace / ".dumont_research"
    prompt_dir = research / "prompts"
    log_dir = research / "logs"
    prompt_dir.mkdir(parents=True, exist_ok=True)
    log_dir.mkdir(parents=True, exist_ok=True)
    prompt = build_prompt(workspace, iteration)
    prompt_path = prompt_dir / f"{iteration:04d}.txt"
    prompt_path.write_text(prompt, encoding="utf-8")
    env = os.environ.copy()
    if args.config:
        env["DUMONT_CONFIG"] = args.config
    command = [
        args.dumont_bin,
        "-p",
        prompt,
        "--model",
        args.model,
        "--output-format",
        "text",
    ]
    started = time.time()
    result = subprocess.run(
        command,
        cwd=workspace,
        env=env,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        timeout=args.timeout,
    )
    elapsed = time.time() - started
    log_path = log_dir / f"{iteration:04d}.log"
    log_path.write_text(result.stdout, encoding="utf-8", errors="replace")
    append_jsonl(
        research / "driver.jsonl",
        {
            "iteration": iteration,
            "exit_code": result.returncode,
            "elapsed_seconds": round(elapsed, 3),
            "log": str(log_path),
            "prompt": str(prompt_path),
            "time": int(time.time()),
        },
    )
    return result.returncode


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workspace", required=True)
    parser.add_argument("--dumont-bin", required=True)
    parser.add_argument("--config", default="")
    parser.add_argument("--model", default="minimax/MiniMax-M3")
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--sleep", type=int, default=30)
    parser.add_argument("--timeout", type=int, default=1800)
    args = parser.parse_args()
    workspace = Path(args.workspace).resolve()
    research = workspace / ".dumont_research"
    (research / "wiki").mkdir(parents=True, exist_ok=True)
    memory = research / "wiki" / "memory.md"
    if not memory.exists():
        memory.write_text("# Research Memory\n\n", encoding="utf-8")
    state_path = research / "state.json"
    state = {"started": int(time.time()), "workspace": str(workspace), "model": args.model}
    write_json(state_path, state)
    exit_code = 0
    first_iteration = next_iteration(research)
    for offset in range(args.iterations):
        iteration = first_iteration + offset
        try:
            code = run_iteration(args, iteration)
            if code != 0:
                exit_code = code
        except subprocess.TimeoutExpired as exc:
            append_jsonl(
                research / "driver.jsonl",
                {
                    "iteration": iteration,
                    "exit_code": 124,
                    "elapsed_seconds": args.timeout,
                    "error": "timeout",
                    "time": int(time.time()),
                },
            )
            exit_code = 124
        if offset + 1 < args.iterations:
            time.sleep(args.sleep)
    state["finished"] = int(time.time())
    state["exit_code"] = exit_code
    write_json(state_path, state)
    return exit_code


if __name__ == "__main__":
    raise SystemExit(main())