solver.py

#!/usr/bin/env python3
"""Background task manager for the physics solver."""

import os
import sys
import json
import time
import shutil
import subprocess
import threading
import urllib.request
from datetime import datetime, timezone

# Solver source location
SOLVER_SRC = "/app/solver"
WORKSPACE_BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workspaces")
# Files that come from the repo — excluded from downloads
SOLVER_SOURCE_FILES = {"run.py", "spawn.py", "outline.md"}
DEFAULT_SOLVER_URL = os.environ.get(
    "SOLVER_REPO_URL",
    "https://github.com/Usami-Mira/CC_Integrater/raw/refs/heads/main/",
)

tasks = {}
tasks_lock = threading.Lock()


def _log(task_id, msg):
    """Append a line to the task's log file AND print to stdout."""
    task = tasks.get(task_id)
    log_path = None
    if task:
        log_path = os.path.join(task["workspace"], "solver.log")
    if log_path:
        try:
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(f"[{_now()}] {msg}\n")
        except Exception:
            pass
    print(f"[{task_id}] {msg}", flush=True)


def _now():
    return datetime.now(timezone.utc).isoformat()


def fetch_solver_files(repo_url):
    """Download solver files from GitHub raw URL into /app/solver/."""
    os.makedirs(SOLVER_SRC, exist_ok=True)
    for name in ("run.py", "spawn.py", "outline.md", "config.json"):
        p = os.path.join(SOLVER_SRC, name)
        if os.path.exists(p):
            os.remove(p)

    print(f"[INIT] Fetching solver files from: {repo_url}", flush=True)
    files = ["run.py", "spawn.py", "outline.md", "config.json"]
    downloaded = []
    errors = []
    for name in files:
        url = f"{repo_url.rstrip('/')}/{name}"
        try:
            with urllib.request.urlopen(url, timeout=30) as resp:
                data = resp.read()
                with open(os.path.join(SOLVER_SRC, name), "wb") as f:
                    f.write(data)
                downloaded.append(name)
                print(f"[INIT] Downloaded {name}", flush=True)
        except Exception as e:
            errors.append(f"{name}: {e}")
            print(f"[INIT] Failed to download {name}: {e}", flush=True)

    if not downloaded:
        print(f"[INIT] CRITICAL: No solver files downloaded. Errors: {errors}", flush=True)
        return False
    if errors:
        print(f"[INIT] Partial download. Missing: {errors}", flush=True)
    return True


def active_count():
    """Caller must hold tasks_lock."""
    return sum(1 for t in tasks.values() if t["status"] in ("pending", "running"))


def submit_task(problem_text, name=None, config_override=None, secret_key=None):
    """Create a new task, return task_id."""
    print("[SUBMIT] submit_task() called", flush=True)
    try:
        task_id = f"task-{int(time.time()*1000)}"
        workspace = os.path.join(WORKSPACE_BASE, task_id)
        print(f"[SUBMIT] Creating workspace: {workspace}", flush=True)
        os.makedirs(workspace, exist_ok=True)

        # Write problem
        with open(os.path.join(workspace, "problem.md"), "w", encoding="utf-8") as f:
            f.write(problem_text)
        print(f"[SUBMIT] Wrote problem.md", flush=True)

        # Copy solver source into workspace
        copied = []
        for fname in ("run.py", "spawn.py", "outline.md", "config.json"):
            src = os.path.join(SOLVER_SRC, fname)
            exists = os.path.exists(src)
            print(f"[SUBMIT]   {fname}: exists={exists} at {src}", flush=True)
            if exists:
                shutil.copy2(src, os.path.join(workspace, fname))
                copied.append(fname)
        print(f"[SUBMIT] Copied: {copied}", flush=True)

        # Apply config override if provided
        if config_override:
            config_path = os.path.join(workspace, "config.json")
            if os.path.exists(config_path):
                with open(config_path, "r", encoding="utf-8") as f:
                    current = json.load(f)
                current.update(config_override)
                with open(config_path, "w", encoding="utf-8") as f:
                    json.dump(current, f, indent=2, ensure_ascii=False)
                print(f"[SUBMIT] Applied config override: {list(config_override.keys())}", flush=True)
            else:
                with open(config_path, "w", encoding="utf-8") as f:
                    json.dump(config_override, f, indent=2, ensure_ascii=False)
                print(f"[SUBMIT] Created config from override", flush=True)

        # Save secret key if provided (private mode)
        is_private = bool(secret_key)
        if is_private:
            with open(os.path.join(workspace, ".secret"), "w", encoding="utf-8") as f:
                f.write(secret_key)

        task = {
            "task_id": task_id,
            "name": name or f"Problem {len(tasks)+1}",
            "status": "pending",
            "phase": "waiting",
            "workspace": workspace,
            "problem_preview": problem_text[:120],
            "problem_length": len(problem_text),
            "created_at": _now(),
            "started_at": None,
            "finished_at": None,
            "error": None,
            "is_private": is_private,
        }
        with tasks_lock:
            tasks[task_id] = task
        print(f"[SUBMIT] Task registered in memory", flush=True)

        # Create log file
        log_path = os.path.join(workspace, "solver.log")
        with open(log_path, "w", encoding="utf-8") as f:
            f.write(f"[{_now()}] Task {task_id} created\n")

        print(f"[SUBMIT] Created task {task_id} -> {workspace}", flush=True)

        # Spawn background thread — run.py manages its own timeout & concurrency
        t = threading.Thread(target=_run_solver, args=(task_id, workspace), daemon=True)
        t.start()
        print(f"[SUBMIT] Thread started", flush=True)
        return task_id
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        print(f"[SUBMIT] ERROR: {e}", flush=True)
        print(f"[SUBMIT] Traceback:\n{tb}", flush=True)
        raise


def _run_solver(task_id, workspace):
    print(f"[{task_id}] _run_solver thread started", flush=True)
    with tasks_lock:
        tasks[task_id].update({"status": "running", "started_at": _now()})

    _log(task_id, f"Starting solver in {workspace}")
    _log(task_id, f"ANTHROPIC_API_KEY: {'set' if os.environ.get('ANTHROPIC_API_KEY') else 'NOT SET'}")
    _log(task_id, f"ANTHROPIC_BASE_URL: {os.environ.get('ANTHROPIC_BASE_URL', '(default)')}")

    env = os.environ.copy()
    start = time.time()
    proc = None

    try:
        cmd = [sys.executable, os.path.join(workspace, "run.py"), workspace]
        print(f"[{task_id}] Running: {' '.join(cmd)}", flush=True)

        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            cwd=workspace,
            env=env,
        )
        _log(task_id, f"spawned run.py with PID {proc.pid}")

        # Read stdout line by line, append to log
        def _read_output():
            if proc.stdout:
                for line in iter(proc.stdout.readline, b""):
                    decoded = line.decode("utf-8", errors="replace").strip()
                    if decoded:
                        _log(task_id, f"[stdout] {decoded}")

        reader = threading.Thread(target=_read_output, daemon=True)
        reader.start()

        # Monitor workspace for progress (no timeout — run.py handles its own)
        last_phase = "starting"
        last_state = _read_state(workspace)
        with tasks_lock:
            tasks[task_id]["file_history"] = []

        while proc.poll() is None:
            phase = _infer_phase(workspace)
            if phase != last_phase:
                _log(task_id, f"Phase: {phase}")
                last_phase = phase

            # Track .state changes for history
            state = _read_state(workspace)
            if state and state != last_state:
                _log(task_id, f"State: {state}")
                history = STATE_TO_HISTORY.get(state, [])
                with tasks_lock:
                    tasks[task_id]["file_history"] = [dict(h) for h in history]
                last_state = state

            with tasks_lock:
                tasks[task_id]["phase"] = phase
                tasks[task_id]["elapsed_s"] = round(time.time() - start)

            time.sleep(10)

        reader.join(timeout=3)
        _log(task_id, f"run.py exited with code {proc.returncode}")

        if proc.returncode != 0:
            _log(task_id, "Solver returned non-zero exit code")
            with tasks_lock:
                tasks[task_id].update({
                    "status": "error",
                    "error": f"Solver exited with code {proc.returncode}",
                    "finished_at": _now(),
                })
            return

        summary = os.path.join(workspace, "final_summary.md")
        if os.path.exists(summary):
            _log(task_id, "Solver completed successfully")
            with tasks_lock:
                tasks[task_id].update({
                    "status": "complete",
                    "phase": "complete",
                    "finished_at": _now(),
                })
        else:
            _log(task_id, "Solver finished but no summary file found")
            with tasks_lock:
                tasks[task_id].update({
                    "status": "error",
                    "error": "Solver finished but no summary file was produced.",
                    "finished_at": _now(),
                })

    except FileNotFoundError:
        _log(task_id, f"FileNotFoundError: run.py not found")
        with tasks_lock:
            tasks[task_id].update({
                "status": "error",
                "error": "run.py not found in workspace.",
                "finished_at": _now(),
            })
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        _log(task_id, f"Exception: {e}")
        _log(task_id, tb)
        print(f"[{task_id}] Traceback:\n{tb}", flush=True)
        with tasks_lock:
            tasks[task_id].update({
                "status": "error",
                "error": str(e),
                "finished_at": _now(),
            })


# Known solver workflow steps in execution order. Unknown files append at end alphabetically.
WORKFLOW_STEPS = ["plan", "solution", "review", "final_summary"]
# Mapping from .state value to completed steps
STATE_TO_HISTORY = {
    "planner": [],  # not started yet
    "builder": [{"step": "plan", "status": "history"}],
    "evaluator": [{"step": "plan", "status": "history"}, {"step": "solution", "status": "history"}],
    "done": [{"step": "plan", "status": "history"}, {"step": "solution", "status": "history"},
             {"step": "review", "status": "history"}, {"step": "final_summary", "status": "history"}],
}
STATE_TO_PHASE = {
    "planner": "starting",
    "builder": "building",
    "evaluator": "evaluating",
    "done": "complete",
}
# Files to exclude from progress tracking
PROGRESS_EXCLUDES = {"problem.md", "solver.log", "outline.md"}


def _read_state(workspace):
    """Read .state file. Returns phase string (planner/builder/evaluator/done) or None."""
    path = os.path.join(workspace, ".state")
    if not os.path.isfile(path):
        return None
    try:
        with open(path, encoding="utf-8") as f:
            return f.read().strip()
    except Exception:
        return None


def _read_secret(workspace):
    """Read .secret file. Returns secret key string or None."""
    path = os.path.join(workspace, ".secret")
    if not os.path.isfile(path):
        return None
    try:
        with open(path, encoding="utf-8") as f:
            return f.read().strip()
    except Exception:
        return None


def _check_access(workspace, provided_key=None):
    """Check if provided_key grants access to a private task. Returns True if public or key matches."""
    secret = _read_secret(workspace)
    if not secret:
        return True  # not a private task
    if provided_key and provided_key == secret:
        return True
    return False


def _build_history_from_state(workspace, state=None):
    """Build progress history from .state file, marking last completed step as 'current'."""
    if state is None:
        state = _read_state(workspace)
    if state is None:
        return None  # no .state → caller should fall back to file scanning

    history = STATE_TO_HISTORY.get(state)
    if history is None:
        return []

    result = list(history)
    if result:
        result[-1] = dict(result[-1], status="current")
    return result


def _check_files(workspace, file_history=None):
    """Return list of steps with history. file_history is an ordered list from monitoring."""
    if file_history and len(file_history) > 0:
        # Use monitored history
        result = list(file_history)
        if result:
            result[-1] = dict(result[-1], status="current")
        return result

    # Priority 1: .state file
    from_state = _build_history_from_state(workspace)
    if from_state is not None:
        return from_state

    # Fallback: scan by mtime (for tasks created before this change or without .state)
    if not os.path.isdir(workspace):
        return []
    mtimes = {}
    for fname in os.listdir(workspace):
        if fname.endswith(".md") and fname not in PROGRESS_EXCLUDES:
            step = fname[:-3]
            if step in WORKFLOW_STEPS:
                mtimes[step] = os.path.getmtime(os.path.join(workspace, fname))
    if not mtimes:
        return []
    latest = max(mtimes, key=mtimes.get)
    result = []
    for step in WORKFLOW_STEPS:
        if step in mtimes:
            status = "current" if step == latest else "history"
            result.append({"step": step, "status": status})
    other = sorted(s for s in mtimes if s not in set(WORKFLOW_STEPS))
    for s in other:
        result.append({"step": s, "status": "history"})
    return result


def _infer_phase(workspace):
    # Priority: .state file first
    state = _read_state(workspace)
    if state and state in STATE_TO_PHASE:
        return STATE_TO_PHASE[state]
    # Fallback: scan .md files
    if os.path.exists(os.path.join(workspace, "final_summary.md")):
        return "complete"
    if os.path.exists(os.path.join(workspace, "review.md")):
        return "evaluating"
    if os.path.exists(os.path.join(workspace, "solution.md")):
        return "building"
    if os.path.exists(os.path.join(workspace, "plan.md")):
        return "planning"
    return "starting"


def get_task(task_id, secret_key=None):
    with tasks_lock:
        task = tasks.get(task_id)
        if task:
            # Check access for private tasks
            if not _check_access(task["workspace"], secret_key):
                return None
            result = dict(task)
            result.pop("problem_preview", None)
            history = result.get("file_history")
            result["files_done"] = _check_files(task["workspace"], history)
            return result
        return None


def get_task_result(task_id, secret_key=None):
    task = get_task(task_id, secret_key)
    if not task:
        return None
    summary = os.path.join(task.get("workspace", ""), "final_summary.md")
    if os.path.exists(summary):
        with open(summary, encoding="utf-8") as f:
            return f.read()
    return None


def get_task_logs(task_id, lines=100, secret_key=None):
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return [f"Task {task_id} not found in memory."]
        if not _check_access(task["workspace"], secret_key):
            return ["Access denied: invalid or missing secret key."]
    log_path = os.path.join(task["workspace"], "solver.log")
    if not os.path.exists(log_path):
        return ["Waiting for solver to start..."]
    try:
        with open(log_path, encoding="utf-8") as f:
            all_lines = f.readlines()
            return [l.rstrip() for l in all_lines[-lines:]]
    except Exception as e:
        return [f"Error reading log: {e}"]


def get_workspace_file(task_id, filename, secret_key=None):
    """Read a file from a task's workspace. Returns content or None if not found."""
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return None
        if not _check_access(task["workspace"], secret_key):
            return None
    fpath = os.path.join(task["workspace"], filename)
    if not os.path.isfile(fpath):
        return None
    with open(fpath, encoding="utf-8") as f:
        return f.read()


def list_workspace_files(task_id, secret_key=None):
    """List all non-hidden .md files in a task's workspace."""
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return []
        if not _check_access(task["workspace"], secret_key):
            return []
    ws = task["workspace"]
    if not os.path.isdir(ws):
        return []
    result = []
    for fname in sorted(os.listdir(ws)):
        if fname.endswith(".md") and not fname.startswith("."):
            result.append(fname)
    return result


def get_running_tasks():
    """Return list of currently running tasks (for frontend recovery after page refresh)."""
    with tasks_lock:
        return [
            {
                "task_id": t["task_id"],
                "status": t["status"],
                "phase": t["phase"],
                "problem": t["problem"],
                "created_at": t["created_at"],
                "files_done": _check_files(t["workspace"], t.get("file_history")),
            }
            for t in tasks.values()
            if t["status"] in ("pending", "running")
        ]


def list_tasks(secret_key=None):
    with tasks_lock:
        return [
            {
                "task_id": t["task_id"],
                "name": t["name"],
                "status": t["status"],
                "phase": t["phase"],
                "problem_preview": t.get("problem_preview", "") if not t.get("is_private") else "🔒",
                "problem_length": t.get("problem_length", 0),
                "created_at": t["created_at"],
                "started_at": t.get("started_at"),
                "finished_at": t.get("finished_at"),
                "files_done": _check_files(t["workspace"], t.get("file_history")),
                "error": t.get("error"),
                "is_private": t.get("is_private", False),
            }
            for t in sorted(tasks.values(), key=lambda x: x["created_at"], reverse=True)
        ]


def delete_task(task_id, secret_key=None):
    """Remove a task from memory and delete its workspace. Returns True if found."""
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return False
        if task.get("is_private") and not _check_access(task["workspace"], secret_key):
            return False
        task = tasks.pop(task_id, None)
    workspace = task.get("workspace", "")
    if os.path.isdir(workspace):
        shutil.rmtree(workspace)
    _log(task_id, f"Task {task_id} deleted")
    return True


def cleanup_completed_tasks():
    """Delete all completed and errored tasks. Returns count of deleted tasks."""
    with tasks_lock:
        to_delete = [
            tid for tid, t in tasks.items()
            if t["status"] in ("complete", "error")
        ]
    count = 0
    for tid in to_delete:
        if delete_task(tid):
            count += 1
    return count


def update_task_config(task_id, config_dict):
    """Update config.json in a task's workspace. Only allowed for pending tasks."""
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return None, "Task not found"
        if task["status"] != "pending":
            return None, f"Cannot update config for task in '{task['status']}' status (only 'pending')"
    config_path = os.path.join(task["workspace"], "config.json")
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            current = json.load(f)
        current.update(config_dict)
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(current, f, indent=2, ensure_ascii=False)
        return current, None
    except FileNotFoundError:
        return None, "config.json not found in workspace"
    except json.JSONDecodeError:
        return None, "Invalid config.json in workspace"