from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Depends
from fastapi.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from sqlalchemy.orm import Session
import database as db
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, StreamingResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
import ipaddress
from typing import Dict, List, Optional, Any, Union, AsyncGenerator, Set
from enum import Enum
import asyncio
import psutil
import os
import json
from datetime import datetime, timedelta
import subprocess
from typing import List, Dict, Optional
import logging
import subprocess
from pathlib import Path
import time
import glob
import re
import markdown
import aiofiles
import tempfile
import uuid
import hashlib
import threading
from typing import Sequence
import sys

# Ensure repo root is importable when running uvicorn from api/ without start script
try:
    _WS_ROOT = Path(__file__).resolve().parents[2]
    if str(_WS_ROOT) not in sys.path:
        sys.path.insert(0, str(_WS_ROOT))
except Exception:
    pass
from agi_dw.core.llm.hf_client import HFClient  # core model interface
from agi_dw.core.planner.service import (
    PlannerConfig,
    VerifierConfig as PlannerVerifierConfig,
    WMConfig,
    ContextAugment,
    plan_with_context,
)
from agi_dw.core.world_model.service import WorldModelService
from agi_dw.core.actuator.code_actions import execute_code_action
from agi_dw.core.actuator.service import (
    select_action as actuator_select_action,
    ActuatorConfig as ActCfg,
    RouterExtras as ActRouterExtras,
    RouterVerifierConfig as ActRouterVerifierCfg,
    WMPriorConfig as ActWMPriorCfg,
    WMScreenConfig as ActWMScreenCfg,
    RepairConfig as ActRepairCfg,
)
from agi_dw.core.verifier.service import VerifierServiceConfig, verify as verifier_run
from agi_dw.core.updater import Updater

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('api_server.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(title="AGI Control Center")

# Initialize database
db.init_db()

# Model registry types
class ModelParams(BaseModel):
    device: Optional[str] = "cuda"
    dtype: Optional[str] = "float16"
    max_memory: Optional[Dict[str, str]] = None
    trust_remote_code: Optional[bool] = False
    use_auth_token: Optional[str] = None
    revision: Optional[str] = None

class ModelRequest(BaseModel):
    model_id: str
    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1
    wm_checkpoint: Optional[str] = None

class ActuatorRequest(BaseModel):
    repo: str
    action: Dict[str, Any]

class ActuatorSelectRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    domain: str = "cli"  # "cli" | "dom"
    mode: str = "router"  # "template" | "two_head" | "router" | "t5" | "nn"
    t5_model: Optional[str] = None
    il_path: Optional[str] = None
    learned_router: bool = False
    router_model_path: Optional[str] = None
    router_threshold: float = 0.5
    router_use_packed_threshold: bool = False
    router_thresholds_json: Optional[str] = None
    dom_structured: bool = False
    verifier_model: Optional[str] = None
    verifier_backend: str = "hf"
    verifier_adapter_dir: Optional[str] = None
    verifier_adapter_bank: Optional[str] = None
    verifier_structured_mode: str = "json"
    verifier_timeout_sec: int = 15
    wm_prior_enabled: bool = False
    wm_prior_model_path: Optional[str] = None
    wm_screen_enabled: bool = False
    wm_screen_threshold: float = 0.7
    repair: bool = True
    prefer_obs_args: bool = False
    default_url: Optional[str] = None
    default_selector: Optional[str] = None

class VerifyRequest(BaseModel):
    trace: Dict[str, Any]
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    backend: str = "hf"
    timeout_sec: int = 15
    adapter_dir: Optional[str] = None
    adapter_bank: Optional[str] = None
    structured_mode: str = "json"
    strict: bool = False
    calibrate: bool = False
    calib_model: Optional[str] = None
    log_prompts: bool = False

class UpdaterRequest(BaseModel):
    fast: bool = False

# Configure CORS for internal network only
# CORS: allow local and same-origin access for browser dashboard
app.add_middleware(
    CORSMiddleware,
    allow_origins=[],
    allow_origin_regex=r"^https?://(localhost|127\\.0\\.0\\.1|10\\.\d+\\.\d+\\.\d+|192\\.168\\.\d+\\.\d+|172\\.(1[6-9]|2\d|3[0-1])\\.\d+\\.\d+)(:\\d+)?$",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup templates and static
BASE_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = BASE_DIR.parent
templates = Jinja2Templates(directory=BASE_DIR / "templates")
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")

# Global state
active_tasks: Dict[str, Dict[str, Any]] = {}
task_processes: Dict[str, asyncio.subprocess.Process] = {}
model_cache: Dict[str, Any] = {}

# Built-in lightweight debug client to enable dev without heavy deps
class DebugClient:
    def __init__(self, model_id: str = "debug/echo") -> None:
        self.model_id = model_id

    def attach_adapter(self, adapter_dir: Optional[str]) -> None:
        return

    def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
        # Very simple behaviors for quick testing in dev
        text = prompt
        if self.model_id == "debug/upper":
            text = prompt.upper()
        elif self.model_id == "debug/repeat":
            text = (prompt + "\n") * max(1, min(3, int(max_new_tokens // max(1, len(prompt)) or 1)))
        else:
            # default echo: return prompt suffix to simulate completion
            text = prompt
        # Apply stop tokens if any
        try:
            if stop:
                cut = len(text)
                for s in stop:
                    if not s:
                        continue
                    idx = text.find(s)
                    if idx != -1:
                        cut = min(cut, idx)
                text = text[:cut]
        except Exception:
            pass
        # Heuristic: return only the completion if prompt looks like chat
        # For debug we simply return the last line after a known separator
        for sep in ["\nassistant:", "\nAssistant:", "\nASSISTANT:"]:
            if sep in text:
                text = text.split(sep, 1)[-1]
                break
        return (text or "").lstrip("\n").rstrip()

    def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
        parts: List[str] = []
        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")
            parts.append(f"{role}: {content}")
        prompt = "\n".join(parts) + "\nassistant:"
        return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)

# Web sessions (in-memory) for model-driven browser
web_sessions: Dict[str, Dict[str, Any]] = {}
web_sessions_lock = threading.Lock()

# ------------- Request ID and Audit Middleware -------------
AUDIT_DIR = (BASE_DIR.parent / "data" / "traces" / "api")
AUDIT_DIR.mkdir(parents=True, exist_ok=True)

class AuditMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        rid = request.headers.get("X-Request-ID") or uuid.uuid4().hex
        start = time.time()
        # process
        response = await call_next(request)
        dur_ms = int((time.time() - start) * 1000)
        try:
            with (AUDIT_DIR / "requests.jsonl").open("a", encoding="utf-8") as f:
                rec = {
                    "ts": datetime.utcnow().isoformat(),
                    "rid": rid,
                    "method": request.method,
                    "path": request.url.path,
                    "status": response.status_code,
                    "duration_ms": dur_ms,
                    "client": request.client.host if request.client else None,
                }
                f.write(json.dumps(rec) + "\n")
        except Exception:
            pass
        response.headers["X-Request-ID"] = rid
        return response

app.add_middleware(AuditMiddleware)

# ---------- Network restriction: IP allowlist middleware ----------
class IPAllowlistMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        # RFC1918 private IPv4 ranges, localhost, and IPv6 loopback + ULA
        self.allowed_networks = [
            ipaddress.ip_network("127.0.0.0/8"),
            ipaddress.ip_network("10.0.0.0/8"),
            ipaddress.ip_network("172.16.0.0/12"),
            ipaddress.ip_network("192.168.0.0/16"),
            ipaddress.ip_network("::1/128"),
            ipaddress.ip_network("fc00::/7"),
        ]

    async def dispatch(self, request, call_next):
        # Determine the peer IP from the connection
        peer_ip = request.client.host if (request.client and request.client.host) else None

        # Only trust X-Forwarded-For if the direct peer is already from an allowed network (trusted proxy)
        xff_header = request.headers.get("x-forwarded-for") or request.headers.get("X-Forwarded-For")
        xff_ip = (xff_header.split(",")[0] or "").strip() if xff_header else None

        def is_allowed(ip_str: str) -> bool:
            try:
                ip_obj = ipaddress.ip_address(ip_str)
                return any(ip_obj in net for net in self.allowed_networks)
            except Exception:
                return False

        # Choose which IP to evaluate: default to the peer IP; if the peer is allowed (proxy), honor XFF
        chosen_ip = None
        if peer_ip:
            chosen_ip = xff_ip if (is_allowed(peer_ip) and xff_ip) else peer_ip

        # Default deny if no usable IP or if chosen IP is not allowed
        if not chosen_ip or not is_allowed(chosen_ip):
            return PlainTextResponse("Forbidden: external access is not allowed", status_code=403)

        return await call_next(request)

# Install the IP allowlist middleware early (after Audit so rejections still log)
app.add_middleware(IPAllowlistMiddleware)

# GPU monitoring functions
def get_gpu_metrics() -> List[Dict[str, Any]]:
    """Get NVIDIA GPU metrics using nvidia-smi"""
    try:
        # Run nvidia-smi with XML output for easier parsing
        result = subprocess.run([
            'nvidia-smi',
            '--query-gpu=index,name,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.used',
            '--format=csv,noheader,nounits'
        ], capture_output=True, text=True, check=True)

        gpus = []
        for line in result.stdout.strip().split('\n'):
            if line:
                idx, name, temp, util, mem_util, mem_total, mem_used = line.split(',')
                gpus.append({
                    'index': int(idx),
                    'name': name.strip(),
                    'temperature': float(temp),
                    'utilization': float(util),
                    'memory_utilization': float(mem_util),
                    'memory_total': float(mem_total),
                    'memory_used': float(mem_used),
                    'timestamp': datetime.now().isoformat()
                })
        return gpus
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.warning("nvidia-smi not available or failed")
        return []

# Resource models
class ResourceRequest(BaseModel):
    cpu_cores: Optional[float] = None  # Number of CPU cores (can be fractional)
    memory_mb: Optional[int] = None    # Memory in MB
    gpu_indices: Optional[List[int]] = None  # List of GPU indices to use
    gpu_memory_mb: Optional[int] = None  # GPU memory per GPU in MB

# Resource tracking
class ResourceManager:
    def __init__(self):
        self.allocated_cpu: Dict[str, float] = {}  # task_id -> cpu_cores
        self.allocated_memory: Dict[str, int] = {}  # task_id -> memory_mb
        self.allocated_gpus: Dict[str, List[int]] = {}  # task_id -> gpu_indices
        self.allocated_gpu_memory: Dict[str, int] = {}  # task_id -> gpu_memory_mb per GPU

        # Get system resources
        self.total_cpu = psutil.cpu_count()
        self.total_memory = psutil.virtual_memory().total // (1024 * 1024)  # Convert to MB
        self.total_gpus = len(get_gpu_metrics())

        if self.total_gpus > 0:
            gpu_info = get_gpu_metrics()[0]  # Use first GPU as reference
            self.gpu_memory = gpu_info["memory_total"]  # Already in MB
        else:
            self.gpu_memory = 0

    def can_allocate(self, request: ResourceRequest) -> bool:
        """Check if requested resources are available"""
        if request is None:
            return True

        # Calculate current allocations
        used_cpu = sum(self.allocated_cpu.values())
        used_memory = sum(self.allocated_memory.values())
        used_gpus = set().union(*self.allocated_gpus.values()) if self.allocated_gpus else set()

        # Check CPU
        if request.cpu_cores and used_cpu + request.cpu_cores > self.total_cpu:
            return False

        # Check Memory
        if request.memory_mb and used_memory + request.memory_mb > self.total_memory:
            return False

        # Check GPUs
        if request.gpu_indices:
            # Check if requested GPUs are available
            if not all(idx < self.total_gpus for idx in request.gpu_indices):
                return False
            if any(idx in used_gpus for idx in request.gpu_indices):
                return False
            # Check GPU memory
            if request.gpu_memory_mb and request.gpu_memory_mb > self.gpu_memory:
                return False

        return True

    def allocate(self, task_id: str, request: ResourceRequest):
        """Allocate resources to a task"""
        if request is None:
            return

        if request.cpu_cores:
            self.allocated_cpu[task_id] = request.cpu_cores
        if request.memory_mb:
            self.allocated_memory[task_id] = request.memory_mb
        if request.gpu_indices:
            self.allocated_gpus[task_id] = request.gpu_indices
            if request.gpu_memory_mb:
                self.allocated_gpu_memory[task_id] = request.gpu_memory_mb

    def release(self, task_id: str):
        """Release resources allocated to a task"""
        self.allocated_cpu.pop(task_id, None)
        self.allocated_memory.pop(task_id, None)
        self.allocated_gpus.pop(task_id, None)
        self.allocated_gpu_memory.pop(task_id, None)

    def get_usage(self) -> Dict[str, Any]:
        """Get current resource usage"""
        return {
            "cpu": {
                "total": self.total_cpu,
                "used": sum(self.allocated_cpu.values()),
                "allocations": self.allocated_cpu
            },
            "memory": {
                "total": self.total_memory,
                "used": sum(self.allocated_memory.values()),
                "allocations": self.allocated_memory
            },
            "gpu": {
                "total": self.total_gpus,
                "used_gpus": list(set().union(*self.allocated_gpus.values())) if self.allocated_gpus else [],
                "allocations": self.allocated_gpus,
                "memory": {
                    "total": self.gpu_memory,
                    "allocations": self.allocated_gpu_memory
                }
            }
        }

resource_manager = ResourceManager()

# Metrics history
metrics_history: List[Dict[str, Any]] = []
gpu_metrics_history: List[Dict[str, Any]] = []
MAX_HISTORY = 100

# Task definitions from Makefile
TASK_CATEGORIES = {
    "benchmarks": [
        "bench.run.humaneval",
        "bench.run.mbpp",
        "bench.run.apps",
        "bench.run.swebench_lite",
        "bench.run.all"
    ],
    "training": [
        "train.sft.plan",
        "train.sft.patch",
        "train.sft.cli",
        "train.sft.policy",
        "train.sft.all"
    ],
    "hitl": [
        "hitl.smoke",
        "hitl.run",
        "hitl.export",
        "hitl.dashboard"
    ]
}

# ---------------- Health ----------------
@app.get("/healthz")
async def healthz() -> Dict[str, str]:
    return {"status": "ok"}

@app.get("/readyz")
async def readyz() -> Dict[str, Any]:
    try:
        # minimal DB probe
        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL
    dependencies: Optional[List[str]] = None  # List of task IDs this task depends on
    resources: Optional[ResourceRequest] = None  # Resource requirements
    ttl_seconds: Optional[int] = None  # Optional TTL for the task

# Note: ModelRequest is already defined near the top with ModelParams; avoid duplicate definitions

async def run_task(task_id: str, request: TaskRequest):
    """Background task to run make targets"""
    try:
        # Get database session
        async_session = SessionLocal()

        # Wait for dependencies to complete
        if request.dependencies:
            while not can_start_task(task_id):
                # Check if any dependency failed
                for dep_id in request.dependencies:
                    if active_tasks[dep_id]["status"] in ["failed", "stopped"]:
                        # Update database
                        db.update_task(async_session, task_id, {
                            "status": db.TaskStatus.FAILED,
                            "error": f"Dependency {dep_id} failed or was stopped"
                        })
                        # Update in-memory state
                        active_tasks[task_id]["status"] = "failed"
                        active_tasks[task_id]["error"] = f"Dependency {dep_id} failed or was stopped"
                        return
                await asyncio.sleep(1)

        # Update database
        db.update_task(async_session, task_id, {
# ... truncated ...


ModelParams

from agi_dw.core.updater import Updater

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('api_server.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(title="AGI Control Center")

# Initialize database
db.init_db()

# Model registry types
class ModelParams(BaseModel):
    device: Optional[str] = "cuda"
    dtype: Optional[str] = "float16"
    max_memory: Optional[Dict[str, str]] = None
    trust_remote_code: Optional[bool] = False
    use_auth_token: Optional[str] = None
    revision: Optional[str] = None

class ModelRequest(BaseModel):
    model_id: str
    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False


ModelRequest

        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(title="AGI Control Center")

# Initialize database
db.init_db()

# Model registry types
class ModelParams(BaseModel):
    device: Optional[str] = "cuda"
    dtype: Optional[str] = "float16"
    max_memory: Optional[Dict[str, str]] = None
    trust_remote_code: Optional[bool] = False
    use_auth_token: Optional[str] = None
    revision: Optional[str] = None

class ModelRequest(BaseModel):
    model_id: str
    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):


InferRequest

# Initialize FastAPI app
app = FastAPI(title="AGI Control Center")

# Initialize database
db.init_db()

# Model registry types
class ModelParams(BaseModel):
    device: Optional[str] = "cuda"
    dtype: Optional[str] = "float16"
    max_memory: Optional[Dict[str, str]] = None
    trust_remote_code: Optional[bool] = False
    use_auth_token: Optional[str] = None
    revision: Optional[str] = None

class ModelRequest(BaseModel):
    model_id: str
    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1


PlannerRequest

# Initialize database
db.init_db()

# Model registry types
class ModelParams(BaseModel):
    device: Optional[str] = "cuda"
    dtype: Optional[str] = "float16"
    max_memory: Optional[Dict[str, str]] = None
    trust_remote_code: Optional[bool] = False
    use_auth_token: Optional[str] = None
    revision: Optional[str] = None

class ModelRequest(BaseModel):
    model_id: str
    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1
    wm_checkpoint: Optional[str] = None

class ActuatorRequest(BaseModel):
    repo: str
    action: Dict[str, Any]

class ActuatorSelectRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    domain: str = "cli"  # "cli" | "dom"
    mode: str = "router"  # "template" | "two_head" | "router" | "t5" | "nn"
    t5_model: Optional[str] = None
    il_path: Optional[str] = None
    learned_router: bool = False


RolloutRequest

    params: Optional[ModelParams] = None

class InferRequest(BaseModel):
    prompt: str
    params: Optional[Dict[str, Any]] = None

class PlannerRequest(BaseModel):
    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1
    wm_checkpoint: Optional[str] = None

class ActuatorRequest(BaseModel):
    repo: str
    action: Dict[str, Any]

class ActuatorSelectRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    domain: str = "cli"  # "cli" | "dom"
    mode: str = "router"  # "template" | "two_head" | "router" | "t5" | "nn"
    t5_model: Optional[str] = None
    il_path: Optional[str] = None
    learned_router: bool = False
    router_model_path: Optional[str] = None
    router_threshold: float = 0.5
    router_use_packed_threshold: bool = False
    router_thresholds_json: Optional[str] = None
    dom_structured: bool = False
    verifier_model: Optional[str] = None
    verifier_backend: str = "hf"


ActuatorRequest

    obs: Dict[str, Any]
    domain: str = "cli"
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    structured: str = "json"
    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1
    wm_checkpoint: Optional[str] = None

class ActuatorRequest(BaseModel):
    repo: str
    action: Dict[str, Any]

class ActuatorSelectRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    domain: str = "cli"  # "cli" | "dom"
    mode: str = "router"  # "template" | "two_head" | "router" | "t5" | "nn"
    t5_model: Optional[str] = None
    il_path: Optional[str] = None
    learned_router: bool = False
    router_model_path: Optional[str] = None
    router_threshold: float = 0.5
    router_use_packed_threshold: bool = False
    router_thresholds_json: Optional[str] = None
    dom_structured: bool = False
    verifier_model: Optional[str] = None
    verifier_backend: str = "hf"
    verifier_adapter_dir: Optional[str] = None
    verifier_adapter_bank: Optional[str] = None
    verifier_structured_mode: str = "json"
    verifier_timeout_sec: int = 15


ActuatorSelectRequest

    candidates: int = 1
    use_tot: bool = False
    verifier_model: Optional[str] = None
    wm_path: Optional[str] = None
    wm_horizon: int = 1
    use_memory: bool = False
    index_k: int = 0
    log_prompts: bool = False

class RolloutRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    actions: List[Dict[str, Any]]
    horizon: int = 1
    wm_checkpoint: Optional[str] = None

class ActuatorRequest(BaseModel):
    repo: str
    action: Dict[str, Any]

class ActuatorSelectRequest(BaseModel):
    obs: Dict[str, Any]
    plan: Dict[str, Any]
    domain: str = "cli"  # "cli" | "dom"
    mode: str = "router"  # "template" | "two_head" | "router" | "t5" | "nn"
    t5_model: Optional[str] = None
    il_path: Optional[str] = None
    learned_router: bool = False
    router_model_path: Optional[str] = None
    router_threshold: float = 0.5
    router_use_packed_threshold: bool = False
    router_thresholds_json: Optional[str] = None
    dom_structured: bool = False
    verifier_model: Optional[str] = None
    verifier_backend: str = "hf"
    verifier_adapter_dir: Optional[str] = None
    verifier_adapter_bank: Optional[str] = None
    verifier_structured_mode: str = "json"
    verifier_timeout_sec: int = 15
    wm_prior_enabled: bool = False
    wm_prior_model_path: Optional[str] = None
    wm_screen_enabled: bool = False
    wm_screen_threshold: float = 0.7
    repair: bool = True
    prefer_obs_args: bool = False
    default_url: Optional[str] = None
    default_selector: Optional[str] = None

class VerifyRequest(BaseModel):
    trace: Dict[str, Any]
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    backend: str = "hf"
    timeout_sec: int = 15
    adapter_dir: Optional[str] = None
    adapter_bank: Optional[str] = None
    structured_mode: str = "json"
    strict: bool = False
    calibrate: bool = False
    calib_model: Optional[str] = None
    log_prompts: bool = False

class UpdaterRequest(BaseModel):
    fast: bool = False

# Configure CORS for internal network only
# CORS: allow local and same-origin access for browser dashboard
app.add_middleware(


VerifyRequest

    router_model_path: Optional[str] = None
    router_threshold: float = 0.5
    router_use_packed_threshold: bool = False
    router_thresholds_json: Optional[str] = None
    dom_structured: bool = False
    verifier_model: Optional[str] = None
    verifier_backend: str = "hf"
    verifier_adapter_dir: Optional[str] = None
    verifier_adapter_bank: Optional[str] = None
    verifier_structured_mode: str = "json"
    verifier_timeout_sec: int = 15
    wm_prior_enabled: bool = False
    wm_prior_model_path: Optional[str] = None
    wm_screen_enabled: bool = False
    wm_screen_threshold: float = 0.7
    repair: bool = True
    prefer_obs_args: bool = False
    default_url: Optional[str] = None
    default_selector: Optional[str] = None

class VerifyRequest(BaseModel):
    trace: Dict[str, Any]
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    backend: str = "hf"
    timeout_sec: int = 15
    adapter_dir: Optional[str] = None
    adapter_bank: Optional[str] = None
    structured_mode: str = "json"
    strict: bool = False
    calibrate: bool = False
    calib_model: Optional[str] = None
    log_prompts: bool = False

class UpdaterRequest(BaseModel):
    fast: bool = False

# Configure CORS for internal network only
# CORS: allow local and same-origin access for browser dashboard
app.add_middleware(
    CORSMiddleware,
    allow_origins=[],
    allow_origin_regex=r"^https?://(localhost|127\\.0\\.0\\.1|10\\.\d+\\.\d+\\.\d+|192\\.168\\.\d+\\.\d+|172\\.(1[6-9]|2\d|3[0-1])\\.\d+\\.\d+)(:\\d+)?$",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup templates and static
BASE_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = BASE_DIR.parent
templates = Jinja2Templates(directory=BASE_DIR / "templates")
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")


UpdaterRequest

    wm_screen_enabled: bool = False
    wm_screen_threshold: float = 0.7
    repair: bool = True
    prefer_obs_args: bool = False
    default_url: Optional[str] = None
    default_selector: Optional[str] = None

class VerifyRequest(BaseModel):
    trace: Dict[str, Any]
    model: str = "meta-llama/Llama-3.1-8B-Instruct"
    backend: str = "hf"
    timeout_sec: int = 15
    adapter_dir: Optional[str] = None
    adapter_bank: Optional[str] = None
    structured_mode: str = "json"
    strict: bool = False
    calibrate: bool = False
    calib_model: Optional[str] = None
    log_prompts: bool = False

class UpdaterRequest(BaseModel):
    fast: bool = False

# Configure CORS for internal network only
# CORS: allow local and same-origin access for browser dashboard
app.add_middleware(
    CORSMiddleware,
    allow_origins=[],
    allow_origin_regex=r"^https?://(localhost|127\\.0\\.0\\.1|10\\.\d+\\.\d+\\.\d+|192\\.168\\.\d+\\.\d+|172\\.(1[6-9]|2\d|3[0-1])\\.\d+\\.\d+)(:\\d+)?$",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup templates and static
BASE_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = BASE_DIR.parent
templates = Jinja2Templates(directory=BASE_DIR / "templates")
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")

# Global state
active_tasks: Dict[str, Dict[str, Any]] = {}


DebugClient

    CORSMiddleware,
    allow_origins=[],
    allow_origin_regex=r"^https?://(localhost|127\\.0\\.0\\.1|10\\.\d+\\.\d+\\.\d+|192\\.168\\.\d+\\.\d+|172\\.(1[6-9]|2\d|3[0-1])\\.\d+\\.\d+)(:\\d+)?$",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup templates and static
BASE_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = BASE_DIR.parent
templates = Jinja2Templates(directory=BASE_DIR / "templates")
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")

# Global state
active_tasks: Dict[str, Dict[str, Any]] = {}
task_processes: Dict[str, asyncio.subprocess.Process] = {}
model_cache: Dict[str, Any] = {}

# Built-in lightweight debug client to enable dev without heavy deps
class DebugClient:
    def __init__(self, model_id: str = "debug/echo") -> None:
        self.model_id = model_id

    def attach_adapter(self, adapter_dir: Optional[str]) -> None:
        return

    def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
        # Very simple behaviors for quick testing in dev
        text = prompt
        if self.model_id == "debug/upper":
            text = prompt.upper()
        elif self.model_id == "debug/repeat":
            text = (prompt + "\n") * max(1, min(3, int(max_new_tokens // max(1, len(prompt)) or 1)))
        else:
            # default echo: return prompt suffix to simulate completion
            text = prompt
        # Apply stop tokens if any
        try:
            if stop:
                cut = len(text)
                for s in stop:
                    if not s:
                        continue
                    idx = text.find(s)
                    if idx != -1:
                        cut = min(cut, idx)
                text = text[:cut]
        except Exception:
            pass
        # Heuristic: return only the completion if prompt looks like chat
        # For debug we simply return the last line after a known separator
        for sep in ["\nassistant:", "\nAssistant:", "\nASSISTANT:"]:
            if sep in text:
                text = text.split(sep, 1)[-1]
                break
        return (text or "").lstrip("\n").rstrip()

    def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
        parts: List[str] = []
        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")
            parts.append(f"{role}: {content}")
        prompt = "\n".join(parts) + "\nassistant:"
        return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)

# Web sessions (in-memory) for model-driven browser
web_sessions: Dict[str, Dict[str, Any]] = {}
web_sessions_lock = threading.Lock()

# ------------- Request ID and Audit Middleware -------------
AUDIT_DIR = (BASE_DIR.parent / "data" / "traces" / "api")
AUDIT_DIR.mkdir(parents=True, exist_ok=True)

class AuditMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        rid = request.headers.get("X-Request-ID") or uuid.uuid4().hex
        start = time.time()
        # process
        response = await call_next(request)
        dur_ms = int((time.time() - start) * 1000)
        try:
            with (AUDIT_DIR / "requests.jsonl").open("a", encoding="utf-8") as f:
                rec = {
                    "ts": datetime.utcnow().isoformat(),


AuditMiddleware

                break
        return (text or "").lstrip("\n").rstrip()

    def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
        parts: List[str] = []
        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")
            parts.append(f"{role}: {content}")
        prompt = "\n".join(parts) + "\nassistant:"
        return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)

# Web sessions (in-memory) for model-driven browser
web_sessions: Dict[str, Dict[str, Any]] = {}
web_sessions_lock = threading.Lock()

# ------------- Request ID and Audit Middleware -------------
AUDIT_DIR = (BASE_DIR.parent / "data" / "traces" / "api")
AUDIT_DIR.mkdir(parents=True, exist_ok=True)

class AuditMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        rid = request.headers.get("X-Request-ID") or uuid.uuid4().hex
        start = time.time()
        # process
        response = await call_next(request)
        dur_ms = int((time.time() - start) * 1000)
        try:
            with (AUDIT_DIR / "requests.jsonl").open("a", encoding="utf-8") as f:
                rec = {
                    "ts": datetime.utcnow().isoformat(),
                    "rid": rid,
                    "method": request.method,
                    "path": request.url.path,
                    "status": response.status_code,
                    "duration_ms": dur_ms,
                    "client": request.client.host if request.client else None,
                }
                f.write(json.dumps(rec) + "\n")
        except Exception:
            pass
        response.headers["X-Request-ID"] = rid
        return response

app.add_middleware(AuditMiddleware)

# ---------- Network restriction: IP allowlist middleware ----------
class IPAllowlistMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        # RFC1918 private IPv4 ranges, localhost, and IPv6 loopback + ULA
        self.allowed_networks = [
            ipaddress.ip_network("127.0.0.0/8"),
            ipaddress.ip_network("10.0.0.0/8"),
            ipaddress.ip_network("172.16.0.0/12"),
            ipaddress.ip_network("192.168.0.0/16"),
            ipaddress.ip_network("::1/128"),
            ipaddress.ip_network("fc00::/7"),
        ]

    async def dispatch(self, request, call_next):
        # Determine the peer IP from the connection
        peer_ip = request.client.host if (request.client and request.client.host) else None


IPAllowlistMiddleware

        try:
            with (AUDIT_DIR / "requests.jsonl").open("a", encoding="utf-8") as f:
                rec = {
                    "ts": datetime.utcnow().isoformat(),
                    "rid": rid,
                    "method": request.method,
                    "path": request.url.path,
                    "status": response.status_code,
                    "duration_ms": dur_ms,
                    "client": request.client.host if request.client else None,
                }
                f.write(json.dumps(rec) + "\n")
        except Exception:
            pass
        response.headers["X-Request-ID"] = rid
        return response

app.add_middleware(AuditMiddleware)

# ---------- Network restriction: IP allowlist middleware ----------
class IPAllowlistMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        # RFC1918 private IPv4 ranges, localhost, and IPv6 loopback + ULA
        self.allowed_networks = [
            ipaddress.ip_network("127.0.0.0/8"),
            ipaddress.ip_network("10.0.0.0/8"),
            ipaddress.ip_network("172.16.0.0/12"),
            ipaddress.ip_network("192.168.0.0/16"),
            ipaddress.ip_network("::1/128"),
            ipaddress.ip_network("fc00::/7"),
        ]

    async def dispatch(self, request, call_next):
        # Determine the peer IP from the connection
        peer_ip = request.client.host if (request.client and request.client.host) else None

        # Only trust X-Forwarded-For if the direct peer is already from an allowed network (trusted proxy)
        xff_header = request.headers.get("x-forwarded-for") or request.headers.get("X-Forwarded-For")
        xff_ip = (xff_header.split(",")[0] or "").strip() if xff_header else None

        def is_allowed(ip_str: str) -> bool:
            try:
                ip_obj = ipaddress.ip_address(ip_str)
                return any(ip_obj in net for net in self.allowed_networks)
            except Exception:
                return False

        # Choose which IP to evaluate: default to the peer IP; if the peer is allowed (proxy), honor XFF
        chosen_ip = None
        if peer_ip:
            chosen_ip = xff_ip if (is_allowed(peer_ip) and xff_ip) else peer_ip

        # Default deny if no usable IP or if chosen IP is not allowed
        if not chosen_ip or not is_allowed(chosen_ip):
            return PlainTextResponse("Forbidden: external access is not allowed", status_code=403)

        return await call_next(request)

# Install the IP allowlist middleware early (after Audit so rejections still log)
app.add_middleware(IPAllowlistMiddleware)

# GPU monitoring functions
def get_gpu_metrics() -> List[Dict[str, Any]]:
    """Get NVIDIA GPU metrics using nvidia-smi"""
    try:
        # Run nvidia-smi with XML output for easier parsing
        result = subprocess.run([
            'nvidia-smi',
            '--query-gpu=index,name,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.used',
            '--format=csv,noheader,nounits'
        ], capture_output=True, text=True, check=True)

        gpus = []
        for line in result.stdout.strip().split('\n'):
            if line:
                idx, name, temp, util, mem_util, mem_total, mem_used = line.split(',')
                gpus.append({


                ip_obj = ipaddress.ip_address(ip_str)
                return any(ip_obj in net for net in self.allowed_networks)
            except Exception:
                return False

        # Choose which IP to evaluate: default to the peer IP; if the peer is allowed (proxy), honor XFF
        chosen_ip = None
        if peer_ip:
            chosen_ip = xff_ip if (is_allowed(peer_ip) and xff_ip) else peer_ip

        # Default deny if no usable IP or if chosen IP is not allowed
        if not chosen_ip or not is_allowed(chosen_ip):
            return PlainTextResponse("Forbidden: external access is not allowed", status_code=403)

        return await call_next(request)

# Install the IP allowlist middleware early (after Audit so rejections still log)
app.add_middleware(IPAllowlistMiddleware)

# GPU monitoring functions
def get_gpu_metrics() -> List[Dict[str, Any]]:
    """Get NVIDIA GPU metrics using nvidia-smi"""
    try:
        # Run nvidia-smi with XML output for easier parsing
        result = subprocess.run([
            'nvidia-smi',
            '--query-gpu=index,name,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.used',
            '--format=csv,noheader,nounits'
        ], capture_output=True, text=True, check=True)

        gpus = []
        for line in result.stdout.strip().split('\n'):
            if line:
                idx, name, temp, util, mem_util, mem_total, mem_used = line.split(',')
                gpus.append({
                    'index': int(idx),
                    'name': name.strip(),
                    'temperature': float(temp),
                    'utilization': float(util),
                    'memory_utilization': float(mem_util),
                    'memory_total': float(mem_total),
                    'memory_used': float(mem_used),
                    'timestamp': datetime.now().isoformat()
                })
        return gpus
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.warning("nvidia-smi not available or failed")
        return []

# Resource models
class ResourceRequest(BaseModel):
    cpu_cores: Optional[float] = None  # Number of CPU cores (can be fractional)
    memory_mb: Optional[int] = None    # Memory in MB
    gpu_indices: Optional[List[int]] = None  # List of GPU indices to use
    gpu_memory_mb: Optional[int] = None  # GPU memory per GPU in MB

# Resource tracking
class ResourceManager:
    def __init__(self):
        self.allocated_cpu: Dict[str, float] = {}  # task_id -> cpu_cores
        self.allocated_memory: Dict[str, int] = {}  # task_id -> memory_mb
        self.allocated_gpus: Dict[str, List[int]] = {}  # task_id -> gpu_indices
        self.allocated_gpu_memory: Dict[str, int] = {}  # task_id -> gpu_memory_mb per GPU

        # Get system resources
        self.total_cpu = psutil.cpu_count()
        self.total_memory = psutil.virtual_memory().total // (1024 * 1024)  # Convert to MB
        self.total_gpus = len(get_gpu_metrics())


ResourceRequest

        gpus = []
        for line in result.stdout.strip().split('\n'):
            if line:
                idx, name, temp, util, mem_util, mem_total, mem_used = line.split(',')
                gpus.append({
                    'index': int(idx),
                    'name': name.strip(),
                    'temperature': float(temp),
                    'utilization': float(util),
                    'memory_utilization': float(mem_util),
                    'memory_total': float(mem_total),
                    'memory_used': float(mem_used),
                    'timestamp': datetime.now().isoformat()
                })
        return gpus
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.warning("nvidia-smi not available or failed")
        return []

# Resource models
class ResourceRequest(BaseModel):
    cpu_cores: Optional[float] = None  # Number of CPU cores (can be fractional)
    memory_mb: Optional[int] = None    # Memory in MB
    gpu_indices: Optional[List[int]] = None  # List of GPU indices to use
    gpu_memory_mb: Optional[int] = None  # GPU memory per GPU in MB

# Resource tracking
class ResourceManager:
    def __init__(self):
        self.allocated_cpu: Dict[str, float] = {}  # task_id -> cpu_cores
        self.allocated_memory: Dict[str, int] = {}  # task_id -> memory_mb
        self.allocated_gpus: Dict[str, List[int]] = {}  # task_id -> gpu_indices
        self.allocated_gpu_memory: Dict[str, int] = {}  # task_id -> gpu_memory_mb per GPU

        # Get system resources
        self.total_cpu = psutil.cpu_count()
        self.total_memory = psutil.virtual_memory().total // (1024 * 1024)  # Convert to MB
        self.total_gpus = len(get_gpu_metrics())

        if self.total_gpus > 0:
            gpu_info = get_gpu_metrics()[0]  # Use first GPU as reference
            self.gpu_memory = gpu_info["memory_total"]  # Already in MB
        else:
            self.gpu_memory = 0


ResourceManager

                    'temperature': float(temp),
                    'utilization': float(util),
                    'memory_utilization': float(mem_util),
                    'memory_total': float(mem_total),
                    'memory_used': float(mem_used),
                    'timestamp': datetime.now().isoformat()
                })
        return gpus
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.warning("nvidia-smi not available or failed")
        return []

# Resource models
class ResourceRequest(BaseModel):
    cpu_cores: Optional[float] = None  # Number of CPU cores (can be fractional)
    memory_mb: Optional[int] = None    # Memory in MB
    gpu_indices: Optional[List[int]] = None  # List of GPU indices to use
    gpu_memory_mb: Optional[int] = None  # GPU memory per GPU in MB

# Resource tracking
class ResourceManager:
    def __init__(self):
        self.allocated_cpu: Dict[str, float] = {}  # task_id -> cpu_cores
        self.allocated_memory: Dict[str, int] = {}  # task_id -> memory_mb
        self.allocated_gpus: Dict[str, List[int]] = {}  # task_id -> gpu_indices
        self.allocated_gpu_memory: Dict[str, int] = {}  # task_id -> gpu_memory_mb per GPU

        # Get system resources
        self.total_cpu = psutil.cpu_count()
        self.total_memory = psutil.virtual_memory().total // (1024 * 1024)  # Convert to MB
        self.total_gpus = len(get_gpu_metrics())

        if self.total_gpus > 0:
            gpu_info = get_gpu_metrics()[0]  # Use first GPU as reference
            self.gpu_memory = gpu_info["memory_total"]  # Already in MB
        else:
            self.gpu_memory = 0

    def can_allocate(self, request: ResourceRequest) -> bool:
        """Check if requested resources are available"""
        if request is None:
            return True

        # Calculate current allocations
        used_cpu = sum(self.allocated_cpu.values())
        used_memory = sum(self.allocated_memory.values())
        used_gpus = set().union(*self.allocated_gpus.values()) if self.allocated_gpus else set()

        # Check CPU
        if request.cpu_cores and used_cpu + request.cpu_cores > self.total_cpu:
            return False

        # Check Memory
        if request.memory_mb and used_memory + request.memory_mb > self.total_memory:
            return False

        # Check GPUs
        if request.gpu_indices:
            # Check if requested GPUs are available
            if not all(idx < self.total_gpus for idx in request.gpu_indices):
                return False
            if any(idx in used_gpus for idx in request.gpu_indices):
                return False
            # Check GPU memory
            if request.gpu_memory_mb and request.gpu_memory_mb > self.gpu_memory:
                return False

        return True

    def allocate(self, task_id: str, request: ResourceRequest):
        """Allocate resources to a task"""
        if request is None:
            return

        if request.cpu_cores:
            self.allocated_cpu[task_id] = request.cpu_cores
        if request.memory_mb:
            self.allocated_memory[task_id] = request.memory_mb
        if request.gpu_indices:
            self.allocated_gpus[task_id] = request.gpu_indices
            if request.gpu_memory_mb:
                self.allocated_gpu_memory[task_id] = request.gpu_memory_mb

    def release(self, task_id: str):
        """Release resources allocated to a task"""
        self.allocated_cpu.pop(task_id, None)
        self.allocated_memory.pop(task_id, None)
        self.allocated_gpus.pop(task_id, None)
        self.allocated_gpu_memory.pop(task_id, None)

    def get_usage(self) -> Dict[str, Any]:
        """Get current resource usage"""
        return {
            "cpu": {
                "total": self.total_cpu,
                "used": sum(self.allocated_cpu.values()),
                "allocations": self.allocated_cpu
            },
            "memory": {
                "total": self.total_memory,
                "used": sum(self.allocated_memory.values()),
                "allocations": self.allocated_memory
            },
            "gpu": {
                "total": self.total_gpus,
                "used_gpus": list(set().union(*self.allocated_gpus.values())) if self.allocated_gpus else [],
                "allocations": self.allocated_gpus,
                "memory": {
                    "total": self.gpu_memory,
                    "allocations": self.allocated_gpu_memory
                }
            }
        }

resource_manager = ResourceManager()

# Metrics history
metrics_history: List[Dict[str, Any]] = []
gpu_metrics_history: List[Dict[str, Any]] = []
MAX_HISTORY = 100

# Task definitions from Makefile
TASK_CATEGORIES = {
    "benchmarks": [
        "bench.run.humaneval",
        "bench.run.mbpp",
        "bench.run.apps",
        "bench.run.swebench_lite",
        "bench.run.all"
    ],
    "training": [
        "train.sft.plan",
        "train.sft.patch",


healthz

        "bench.run.swebench_lite",
        "bench.run.all"
    ],
    "training": [
        "train.sft.plan",
        "train.sft.patch",
        "train.sft.cli",
        "train.sft.policy",
        "train.sft.all"
    ],
    "hitl": [
        "hitl.smoke",
        "hitl.run",
        "hitl.export",
        "hitl.dashboard"
    ]
}

# ---------------- Health ----------------
@app.get("/healthz")
async def healthz() -> Dict[str, str]:
    return {"status": "ok"}

@app.get("/readyz")
async def readyz() -> Dict[str, Any]:
    try:
        # minimal DB probe
        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",


readyz

        "train.sft.plan",
        "train.sft.patch",
        "train.sft.cli",
        "train.sft.policy",
        "train.sft.all"
    ],
    "hitl": [
        "hitl.smoke",
        "hitl.run",
        "hitl.export",
        "hitl.dashboard"
    ]
}

# ---------------- Health ----------------
@app.get("/healthz")
async def healthz() -> Dict[str, str]:
    return {"status": "ok"}

@app.get("/readyz")
async def readyz() -> Dict[str, Any]:
    try:
        # minimal DB probe
        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"


    ]
}

# ---------------- Health ----------------
@app.get("/healthz")
async def healthz() -> Dict[str, str]:
    return {"status": "ok"}

@app.get("/readyz")
async def readyz() -> Dict[str, Any]:
    try:
        # minimal DB probe
        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL


@app.get("/healthz")
async def healthz() -> Dict[str, str]:
    return {"status": "ok"}

@app.get("/readyz")
async def readyz() -> Dict[str, Any]:
    try:
        # minimal DB probe
        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL
    dependencies: Optional[List[str]] = None  # List of task IDs this task depends on
    resources: Optional[ResourceRequest] = None  # Resource requirements
    ttl_seconds: Optional[int] = None  # Optional TTL for the task

# Note: ModelRequest is already defined near the top with ModelParams; avoid duplicate definitions

async def run_task(task_id: str, request: TaskRequest):
    """Background task to run make targets"""
    try:


TaskPriority

        session = db.SessionLocal()
        _ = db.get_system_metrics(session, limit=1)
        session.close()
        return {"ready": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL
    dependencies: Optional[List[str]] = None  # List of task IDs this task depends on
    resources: Optional[ResourceRequest] = None  # Resource requirements
    ttl_seconds: Optional[int] = None  # Optional TTL for the task

# Note: ModelRequest is already defined near the top with ModelParams; avoid duplicate definitions

async def run_task(task_id: str, request: TaskRequest):
    """Background task to run make targets"""
    try:
        # Get database session
        async_session = SessionLocal()

        # Wait for dependencies to complete
        if request.dependencies:
            while not can_start_task(task_id):


TaskRequest

@app.get("/api/version")
async def api_version() -> Dict[str, Any]:
    return {"name": "agi_dw_api", "version": "0.1.0"}

@app.get("/api/info")
async def api_info() -> Dict[str, Any]:
    return {
        "name": "AGI Control Center",
        "version": "0.1.0",
        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL
    dependencies: Optional[List[str]] = None  # List of task IDs this task depends on
    resources: Optional[ResourceRequest] = None  # Resource requirements
    ttl_seconds: Optional[int] = None  # Optional TTL for the task

# Note: ModelRequest is already defined near the top with ModelParams; avoid duplicate definitions

async def run_task(task_id: str, request: TaskRequest):
    """Background task to run make targets"""
    try:
        # Get database session
        async_session = SessionLocal()

        # Wait for dependencies to complete
        if request.dependencies:
            while not can_start_task(task_id):
                # Check if any dependency failed
                for dep_id in request.dependencies:
                    if active_tasks[dep_id]["status"] in ["failed", "stopped"]:
                        # Update database
                        db.update_task(async_session, task_id, {
                            "status": db.TaskStatus.FAILED,
                            "error": f"Dependency {dep_id} failed or was stopped"
                        })


        "workspace": str(WORKSPACE_DIR),
        "time": datetime.now().isoformat(),
    }

class TaskPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"

class TaskRequest(BaseModel):
    task: str
    args: Optional[Dict[str, Any]] = None
    priority: TaskPriority = TaskPriority.NORMAL
    dependencies: Optional[List[str]] = None  # List of task IDs this task depends on
    resources: Optional[ResourceRequest] = None  # Resource requirements
    ttl_seconds: Optional[int] = None  # Optional TTL for the task

# Note: ModelRequest is already defined near the top with ModelParams; avoid duplicate definitions

async def run_task(task_id: str, request: TaskRequest):
    """Background task to run make targets"""
    try:
        # Get database session
        async_session = SessionLocal()

        # Wait for dependencies to complete
        if request.dependencies:
            while not can_start_task(task_id):
                # Check if any dependency failed
                for dep_id in request.dependencies:
                    if active_tasks[dep_id]["status"] in ["failed", "stopped"]:
                        # Update database
                        db.update_task(async_session, task_id, {
                            "status": db.TaskStatus.FAILED,
                            "error": f"Dependency {dep_id} failed or was stopped"
                        })
                        # Update in-memory state
                        active_tasks[task_id]["status"] = "failed"
                        active_tasks[task_id]["error"] = f"Dependency {dep_id} failed or was stopped"
                        return
                await asyncio.sleep(1)

        # Update database
        db.update_task(async_session, task_id, {
            "status": db.TaskStatus.RUNNING,
            "started_at": datetime.now()
        })
        # Update in-memory state
        active_tasks[task_id]["status"] = "running"
        # init live buffers
        active_tasks[task_id]["stdout"] = ""
        active_tasks[task_id]["stderr"] = ""

        # Prepare command
        cmd = f"PYTHONPATH={WORKSPACE_DIR} make -C {WORKSPACE_DIR} {request.task}"
        if request.args:
            args_str = " ".join([f"{k}={v}" for k,v in request.args.items()])
            cmd = f"{cmd} {args_str}"

        # Execute task
        proc = await asyncio.create_subprocess_shell(
            cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        task_processes[task_id] = proc
        # Live log readers
        async def _read_stream(stream: asyncio.StreamReader, key: str) -> None:
            try:
                while True:
                    chunk = await stream.readline()
                    if not chunk:
                        break
                    text = chunk.decode(errors="replace")
                    buf = active_tasks.get(task_id, {}).get(key, "") + text
                    # cap to last 20000 chars
                    if len(buf) > 20000:
                        buf = buf[-20000:]
                    if task_id in active_tasks:
                        active_tasks[task_id][key] = buf
            except Exception:
                pass
        reader_out = asyncio.create_task(_read_stream(proc.stdout, "stdout"))
        reader_err = asyncio.create_task(_read_stream(proc.stderr, "stderr"))

        await proc.wait()
        try:
            await asyncio.wait_for(reader_out, timeout=1.0)
        except Exception:
            reader_out.cancel()
        try:
            await asyncio.wait_for(reader_err, timeout=1.0)
        except Exception:
            reader_err.cancel()

        # Update task results
        # Update database
        status = db.TaskStatus.COMPLETED if proc.returncode == 0 else db.TaskStatus.FAILED
        db.update_task(async_session, task_id, {
            "status": status,
            "completed_at": datetime.now(),
            "return_code": proc.returncode,
            "stdout": active_tasks[task_id].get("stdout", ""),
            "stderr": active_tasks[task_id].get("stderr", "")
        })

        # Update in-memory state
        active_tasks[task_id].update({
            "status": status.value,
            "completed_at": datetime.now().isoformat(),
            "return_code": proc.returncode,
            # stdout/stderr already populated incrementally
        })

        # Release resources on completion
        resource_manager.release(task_id)

    except Exception as e:
        logger.error(f"Error in task {task_id}: {str(e)}")
        active_tasks[task_id].update({
            "status": "failed",
            "error": str(e),
            "completed_at": datetime.now().isoformat()
        })
    finally:
        task_processes.pop(task_id, None)

@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render main dashboard"""
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)


dashboard

            "status": status.value,
            "completed_at": datetime.now().isoformat(),
            "return_code": proc.returncode,
            # stdout/stderr already populated incrementally
        })

        # Release resources on completion
        resource_manager.release(task_id)

    except Exception as e:
        logger.error(f"Error in task {task_id}: {str(e)}")
        active_tasks[task_id].update({
            "status": "failed",
            "error": str(e),
            "completed_at": datetime.now().isoformat()
        })
    finally:
        task_processes.pop(task_id, None)

@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render main dashboard"""
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)
async def selfplay_page(request: Request):
    """Render selfplay control page"""
    return templates.TemplateResponse("selfplay.html", {"request": request})

@app.post("/api/task/start")
async def start_task(
    request: TaskRequest,
    background_tasks: BackgroundTasks,
    database: Session = Depends(db.get_db)


    except Exception as e:
        logger.error(f"Error in task {task_id}: {str(e)}")
        active_tasks[task_id].update({
            "status": "failed",
            "error": str(e),
            "completed_at": datetime.now().isoformat()
        })
    finally:
        task_processes.pop(task_id, None)

@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render main dashboard"""
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)
async def selfplay_page(request: Request):
    """Render selfplay control page"""
    return templates.TemplateResponse("selfplay.html", {"request": request})

@app.post("/api/task/start")
async def start_task(
    request: TaskRequest,
    background_tasks: BackgroundTasks,
    database: Session = Depends(db.get_db)
):
    """Start a make task"""
    # Validate task exists in categories
    valid_tasks = [task for tasks in TASK_CATEGORIES.values() for task in tasks]
    if request.task not in valid_tasks:


            "completed_at": datetime.now().isoformat()
        })
    finally:
        task_processes.pop(task_id, None)

@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render main dashboard"""
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)
async def selfplay_page(request: Request):
    """Render selfplay control page"""
    return templates.TemplateResponse("selfplay.html", {"request": request})

@app.post("/api/task/start")
async def start_task(
    request: TaskRequest,
    background_tasks: BackgroundTasks,
    database: Session = Depends(db.get_db)
):
    """Start a make task"""
    # Validate task exists in categories
    valid_tasks = [task for tasks in TASK_CATEGORIES.values() for task in tasks]
    if request.task not in valid_tasks:
        raise HTTPException(status_code=400, detail="Invalid task")

    task_id = f"{request.task}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"


@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render main dashboard"""
    return templates.TemplateResponse("dashboard.html", {
        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)
async def selfplay_page(request: Request):
    """Render selfplay control page"""
    return templates.TemplateResponse("selfplay.html", {"request": request})

@app.post("/api/task/start")
async def start_task(
    request: TaskRequest,
    background_tasks: BackgroundTasks,
    database: Session = Depends(db.get_db)
):
    """Start a make task"""
    # Validate task exists in categories
    valid_tasks = [task for tasks in TASK_CATEGORIES.values() for task in tasks]
    if request.task not in valid_tasks:
        raise HTTPException(status_code=400, detail="Invalid task")

    task_id = f"{request.task}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    # Validate dependencies if provided
    if request.dependencies:
        # Check if all dependencies exist
        for dep_id in request.dependencies:
            if dep_id not in active_tasks:


        "request": request,
        "tasks": list(active_tasks.values()),
        "categories": TASK_CATEGORIES
    })

@app.get("/dev", response_class=HTMLResponse)
async def dev_workbench(request: Request):
    """Render developer workbench (placeholder)"""
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/dev", response_class=HTMLResponse)
async def dev_page(request: Request):
    return templates.TemplateResponse("dev.html", {"request": request})

@app.get("/selfplay", response_class=HTMLResponse)
async def selfplay_page(request: Request):
    """Render selfplay control page"""
    return templates.TemplateResponse("selfplay.html", {"request": request})

@app.post("/api/task/start")
async def start_task(
    request: TaskRequest,
    background_tasks: BackgroundTasks,
    database: Session = Depends(db.get_db)
):
    """Start a make task"""
    # Validate task exists in categories
    valid_tasks = [task for tasks in TASK_CATEGORIES.values() for task in tasks]
    if request.task not in valid_tasks:
        raise HTTPException(status_code=400, detail="Invalid task")

    task_id = f"{request.task}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    # Validate dependencies if provided
    if request.dependencies:
        # Check if all dependencies exist
        for dep_id in request.dependencies:
            if dep_id not in active_tasks:
                raise HTTPException(status_code=400, detail=f"Dependency task {dep_id} not found")

        # Check for cycles
        if detect_cycles(task_id, request.dependencies):
            raise HTTPException(status_code=400, detail="Circular dependency detected")

        # Get all dependencies (direct and indirect)
        all_deps = get_all_dependencies(task_id, set(request.dependencies))

        # Check for any failed dependencies
        for dep_id in all_deps:
            if active_tasks[dep_id]["status"] == "failed":
                raise HTTPException(status_code=400, detail=f"Dependency task {dep_id} has failed")
            elif active_tasks[dep_id]["status"] == "stopped":
                raise HTTPException(status_code=400, detail=f"Dependency task {dep_id} was stopped")

    # Validate and allocate resources
    if request.resources:
        if not resource_manager.can_allocate(request.resources):
            raise HTTPException(status_code=400, detail="Requested resources are not available")
        resource_manager.allocate(task_id, request.resources)

    # Create task in database
    task_data = {
        "id": task_id,
        "name": request.task,
        "args": request.args,
        "status": db.TaskStatus.QUEUED,
        "progress": 0,
        "priority": request.priority,
        "dependencies": request.dependencies or [],
        "resources": request.resources.dict() if request.resources else None,
        "created_at": datetime.now()
    }

    db_task = db.create_task(database, task_data)

    # Update in-memory state
    active_tasks[task_id] = db_task.to_dict()
    # Attach TTL if provided
    if request.ttl_seconds:
        try:
            active_tasks[task_id]["ttl_seconds"] = int(request.ttl_seconds)
            active_tasks[task_id]["created_at_iso"] = datetime.now().isoformat()
        except Exception:
            pass

    background_tasks.add_task(run_task, task_id, request)

    return {"task_id": task_id}

@app.post("/api/task/stop/{task_id}")
async def stop_task(task_id: str):
    """Stop a running task"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    if task["status"] not in ("running", "queued"):
        raise HTTPException(status_code=400, detail="Task is not running")

    proc = task_processes.get(task_id)
    if proc:
        try:
            proc.terminate()
            try:
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                proc.kill()
        except ProcessLookupError:


        "created_at": datetime.now()
    }

    db_task = db.create_task(database, task_data)

    # Update in-memory state
    active_tasks[task_id] = db_task.to_dict()
    # Attach TTL if provided
    if request.ttl_seconds:
        try:
            active_tasks[task_id]["ttl_seconds"] = int(request.ttl_seconds)
            active_tasks[task_id]["created_at_iso"] = datetime.now().isoformat()
        except Exception:
            pass

    background_tasks.add_task(run_task, task_id, request)

    return {"task_id": task_id}

@app.post("/api/task/stop/{task_id}")
async def stop_task(task_id: str):
    """Stop a running task"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    if task["status"] not in ("running", "queued"):
        raise HTTPException(status_code=400, detail="Task is not running")

    proc = task_processes.get(task_id)
    if proc:
        try:
            proc.terminate()
            try:
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                proc.kill()
        except ProcessLookupError:
            pass

    task["status"] = "stopped"
    task["completed_at"] = datetime.now().isoformat()
    task_processes.pop(task_id, None)

    # Release allocated resources
    resource_manager.release(task_id)

    return {"ok": True}

def _priority_value(priority: str) -> int:
    """Helper to convert priority string to numeric value for sorting"""
    priority_map = {
        TaskPriority.CRITICAL: 0,
        TaskPriority.HIGH: 1,
        TaskPriority.NORMAL: 2,
        TaskPriority.LOW: 3
    }
    return priority_map.get(priority, 999)  # Unknown priorities sort last

def detect_cycles(task_id: str, dependencies: List[str], visited: Set[str] = None, path: Set[str] = None) -> bool:
    """
    Detect cycles in task dependencies using DFS
    Returns True if a cycle is detected, False otherwise
    """
    if visited is None:
        visited = set()
    if path is None:
        path = set()


    proc = task_processes.get(task_id)
    if proc:
        try:
            proc.terminate()
            try:
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                proc.kill()
        except ProcessLookupError:
            pass

    task["status"] = "stopped"
    task["completed_at"] = datetime.now().isoformat()
    task_processes.pop(task_id, None)

    # Release allocated resources
    resource_manager.release(task_id)

    return {"ok": True}

def _priority_value(priority: str) -> int:
    """Helper to convert priority string to numeric value for sorting"""
    priority_map = {
        TaskPriority.CRITICAL: 0,
        TaskPriority.HIGH: 1,
        TaskPriority.NORMAL: 2,
        TaskPriority.LOW: 3
    }
    return priority_map.get(priority, 999)  # Unknown priorities sort last

def detect_cycles(task_id: str, dependencies: List[str], visited: Set[str] = None, path: Set[str] = None) -> bool:
    """
    Detect cycles in task dependencies using DFS
    Returns True if a cycle is detected, False otherwise
    """
    if visited is None:
        visited = set()
    if path is None:
        path = set()

    visited.add(task_id)
    path.add(task_id)

    task_deps = dependencies if task_id not in active_tasks else active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in visited:
            if detect_cycles(dep, dependencies, visited, path):
                return True
        elif dep in path:


    task["status"] = "stopped"
    task["completed_at"] = datetime.now().isoformat()
    task_processes.pop(task_id, None)

    # Release allocated resources
    resource_manager.release(task_id)

    return {"ok": True}

def _priority_value(priority: str) -> int:
    """Helper to convert priority string to numeric value for sorting"""
    priority_map = {
        TaskPriority.CRITICAL: 0,
        TaskPriority.HIGH: 1,
        TaskPriority.NORMAL: 2,
        TaskPriority.LOW: 3
    }
    return priority_map.get(priority, 999)  # Unknown priorities sort last

def detect_cycles(task_id: str, dependencies: List[str], visited: Set[str] = None, path: Set[str] = None) -> bool:
    """
    Detect cycles in task dependencies using DFS
    Returns True if a cycle is detected, False otherwise
    """
    if visited is None:
        visited = set()
    if path is None:
        path = set()

    visited.add(task_id)
    path.add(task_id)

    task_deps = dependencies if task_id not in active_tasks else active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in visited:
            if detect_cycles(dep, dependencies, visited, path):
                return True
        elif dep in path:
            return True

    path.remove(task_id)
    return False

def get_all_dependencies(task_id: str, deps: Set[str] = None) -> Set[str]:
    """
    Get all dependencies for a task (direct and indirect)
    """
    if deps is None:
        deps = set()

    if task_id not in active_tasks:
        return deps

    task_deps = active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in deps:
            deps.add(dep)
            deps.update(get_all_dependencies(dep))

    return deps

def can_start_task(task_id: str) -> bool:


    """
    if visited is None:
        visited = set()
    if path is None:
        path = set()

    visited.add(task_id)
    path.add(task_id)

    task_deps = dependencies if task_id not in active_tasks else active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in visited:
            if detect_cycles(dep, dependencies, visited, path):
                return True
        elif dep in path:
            return True

    path.remove(task_id)
    return False

def get_all_dependencies(task_id: str, deps: Set[str] = None) -> Set[str]:
    """
    Get all dependencies for a task (direct and indirect)
    """
    if deps is None:
        deps = set()

    if task_id not in active_tasks:
        return deps

    task_deps = active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in deps:
            deps.add(dep)
            deps.update(get_all_dependencies(dep))

    return deps

def can_start_task(task_id: str) -> bool:
    """
    Check if a task can be started based on its dependencies
    """
    if task_id not in active_tasks:
        return False

    task = active_tasks[task_id]
    if not task.get("dependencies"):
        return True

    # Check if all dependencies are completed
    for dep_id in task["dependencies"]:
        if dep_id not in active_tasks:
            return False
        dep_status = active_tasks[dep_id]["status"]
        if dep_status != "completed":
            return False


    return False

def get_all_dependencies(task_id: str, deps: Set[str] = None) -> Set[str]:
    """
    Get all dependencies for a task (direct and indirect)
    """
    if deps is None:
        deps = set()

    if task_id not in active_tasks:
        return deps

    task_deps = active_tasks[task_id].get("dependencies", [])
    for dep in task_deps:
        if dep not in deps:
            deps.add(dep)
            deps.update(get_all_dependencies(dep))

    return deps

def can_start_task(task_id: str) -> bool:
    """
    Check if a task can be started based on its dependencies
    """
    if task_id not in active_tasks:
        return False

    task = active_tasks[task_id]
    if not task.get("dependencies"):
        return True

    # Check if all dependencies are completed
    for dep_id in task["dependencies"]:
        if dep_id not in active_tasks:
            return False
        dep_status = active_tasks[dep_id]["status"]
        if dep_status != "completed":
            return False

    return True

@app.get("/api/task/list")
async def list_tasks(
    status: Optional[str] = None,
    priority: Optional[TaskPriority] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100,
    database: Session = Depends(db.get_db)
):
    """List tasks with filters, sorting and pagination"""
    tasks, total = db.get_tasks(
        database,
        status=status,
        priority=priority,
        search=search,
        sort_by=sort_by,
        sort_desc=sort_desc,


    Check if a task can be started based on its dependencies
    """
    if task_id not in active_tasks:
        return False

    task = active_tasks[task_id]
    if not task.get("dependencies"):
        return True

    # Check if all dependencies are completed
    for dep_id in task["dependencies"]:
        if dep_id not in active_tasks:
            return False
        dep_status = active_tasks[dep_id]["status"]
        if dep_status != "completed":
            return False

    return True

@app.get("/api/task/list")
async def list_tasks(
    status: Optional[str] = None,
    priority: Optional[TaskPriority] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100,
    database: Session = Depends(db.get_db)
):
    """List tasks with filters, sorting and pagination"""
    tasks, total = db.get_tasks(
        database,
        status=status,
        priority=priority,
        search=search,
        sort_by=sort_by,
        sort_desc=sort_desc,
        offset=offset,
        limit=limit
    )
    return {
        "items": [task.to_dict() for task in tasks],
        "total": total,
        "offset": offset,
        "limit": limit
    }

@app.post("/api/task/cancel_all")
async def cancel_all_tasks():
    ids = list(active_tasks.keys())
    for tid in ids:
        try:
            await stop_task(tid)
        except Exception:
            continue
    return {"ok": True, "cancelled": ids}

@app.get("/api/task/dependencies/{task_id}")
async def get_task_dependencies(task_id: str):
    """Get task dependencies and their status"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    direct_deps = task.get("dependencies", [])
    all_deps = get_all_dependencies(task_id)


):
    """List tasks with filters, sorting and pagination"""
    tasks, total = db.get_tasks(
        database,
        status=status,
        priority=priority,
        search=search,
        sort_by=sort_by,
        sort_desc=sort_desc,
        offset=offset,
        limit=limit
    )
    return {
        "items": [task.to_dict() for task in tasks],
        "total": total,
        "offset": offset,
        "limit": limit
    }

@app.post("/api/task/cancel_all")
async def cancel_all_tasks():
    ids = list(active_tasks.keys())
    for tid in ids:
        try:
            await stop_task(tid)
        except Exception:
            continue
    return {"ok": True, "cancelled": ids}

@app.get("/api/task/dependencies/{task_id}")
async def get_task_dependencies(task_id: str):
    """Get task dependencies and their status"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    direct_deps = task.get("dependencies", [])
    all_deps = get_all_dependencies(task_id)

    # Build dependency tree
    dep_tree = {
        "task_id": task_id,
        "status": task["status"],
        "direct_dependencies": [
            {
                "task_id": dep_id,
                "status": active_tasks[dep_id]["status"] if dep_id in active_tasks else "unknown"
            }


        limit=limit
    )
    return {
        "items": [task.to_dict() for task in tasks],
        "total": total,
        "offset": offset,
        "limit": limit
    }

@app.post("/api/task/cancel_all")
async def cancel_all_tasks():
    ids = list(active_tasks.keys())
    for tid in ids:
        try:
            await stop_task(tid)
        except Exception:
            continue
    return {"ok": True, "cancelled": ids}

@app.get("/api/task/dependencies/{task_id}")
async def get_task_dependencies(task_id: str):
    """Get task dependencies and their status"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    direct_deps = task.get("dependencies", [])
    all_deps = get_all_dependencies(task_id)

    # Build dependency tree
    dep_tree = {
        "task_id": task_id,
        "status": task["status"],
        "direct_dependencies": [
            {
                "task_id": dep_id,
                "status": active_tasks[dep_id]["status"] if dep_id in active_tasks else "unknown"
            }
            for dep_id in direct_deps
        ],
        "all_dependencies": [
            {
                "task_id": dep_id,
                "status": active_tasks[dep_id]["status"] if dep_id in active_tasks else "unknown"
            }
            for dep_id in all_deps
        ],
        "can_start": can_start_task(task_id)
    }

    return dep_tree

@app.get("/api/task/logs/{task_id}")
async def get_task_logs(task_id: str):
    """Get logs for a task"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    return {
        "stdout": task.get("stdout", ""),
        "stderr": task.get("stderr", "")
    }

@app.get("/api/metrics")
async def get_metrics(database: Session = Depends(db.get_db)):
    """Get system metrics history"""
    now = time.time()

    # Get GPU metrics
    gpu_metrics = get_gpu_metrics()


        "direct_dependencies": [
            {
                "task_id": dep_id,
                "status": active_tasks[dep_id]["status"] if dep_id in active_tasks else "unknown"
            }
            for dep_id in direct_deps
        ],
        "all_dependencies": [
            {
                "task_id": dep_id,
                "status": active_tasks[dep_id]["status"] if dep_id in active_tasks else "unknown"
            }
            for dep_id in all_deps
        ],
        "can_start": can_start_task(task_id)
    }

    return dep_tree

@app.get("/api/task/logs/{task_id}")
async def get_task_logs(task_id: str):
    """Get logs for a task"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    return {
        "stdout": task.get("stdout", ""),
        "stderr": task.get("stderr", "")
    }

@app.get("/api/metrics")
async def get_metrics(database: Session = Depends(db.get_db)):
    """Get system metrics history"""
    now = time.time()

    # Get GPU metrics
    gpu_metrics = get_gpu_metrics()

    # Add system metrics to database
    metrics_data = {
        "timestamp": datetime.now(),
        "cpu_percent": psutil.cpu_percent(),
        "memory_percent": psutil.virtual_memory().percent,
        "active_tasks": len([t for t in active_tasks.values() if t["status"] == "running"]),
        "gpu_metrics": gpu_metrics
    }
    db.add_system_metrics(database, metrics_data)

    # Get recent metrics from database


            for dep_id in all_deps
        ],
        "can_start": can_start_task(task_id)
    }

    return dep_tree

@app.get("/api/task/logs/{task_id}")
async def get_task_logs(task_id: str):
    """Get logs for a task"""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = active_tasks[task_id]
    return {
        "stdout": task.get("stdout", ""),
        "stderr": task.get("stderr", "")
    }

@app.get("/api/metrics")
async def get_metrics(database: Session = Depends(db.get_db)):
    """Get system metrics history"""
    now = time.time()

    # Get GPU metrics
    gpu_metrics = get_gpu_metrics()

    # Add system metrics to database
    metrics_data = {
        "timestamp": datetime.now(),
        "cpu_percent": psutil.cpu_percent(),
        "memory_percent": psutil.virtual_memory().percent,
        "active_tasks": len([t for t in active_tasks.values() if t["status"] == "running"]),
        "gpu_metrics": gpu_metrics
    }
    db.add_system_metrics(database, metrics_data)

    # Get recent metrics from database
    metrics = db.get_system_metrics(database, limit=MAX_HISTORY)
    metrics_list = [m.to_dict() for m in metrics]

    # Prepare GPU metrics series
    gpu_series = {}
    if any(m.get("gpu_metrics") for m in metrics_list):
        gpu_indices = {gpu["index"] for m in metrics_list if m.get("gpu_metrics") for gpu in m["gpu_metrics"]}
        for idx in gpu_indices:
            gpu_series[f"gpu_{idx}_util"] = []
            gpu_series[f"gpu_{idx}_mem"] = []
            gpu_series[f"gpu_{idx}_temp"] = []

        for m in metrics_list:
            gpus = {gpu["index"]: gpu for gpu in m.get("gpu_metrics", [])}
            for idx in gpu_indices:
                gpu = gpus.get(idx, {})
                gpu_series[f"gpu_{idx}_util"].append(gpu.get("utilization"))
                gpu_series[f"gpu_{idx}_mem"].append(gpu.get("memory_utilization"))
                gpu_series[f"gpu_{idx}_temp"].append(gpu.get("temperature"))

    return {
        "timestamps": [m["timestamp"] for m in metrics_list],
        "cpu_usage": [m["cpu_percent"] for m in metrics_list],
        "memory_usage": [m["memory_percent"] for m in metrics_list],
        "active_tasks": [m["active_tasks"] for m in metrics_list],
        "gpu_metrics": gpu_series
    }

@app.get("/api/gpu/info")
async def get_gpu_info():
    """Get current GPU information"""
    return {"gpus": get_gpu_metrics()}

# ---------------- Scripts listing ----------------
@app.get("/api/scripts/list")
async def scripts_list() -> Dict[str, Any]:
    root = WORKSPACE_DIR / "scripts"
    items: List[Dict[str, Any]] = []
    if root.exists():
        for p in sorted(root.rglob("*.py")):
            if p.is_file() and not p.name.startswith('.'):
                try:
                    items.append({
                        "path": str(p.relative_to(WORKSPACE_DIR)),
                        "size": p.stat().st_size
                    })
                except Exception:


            gpu_series[f"gpu_{idx}_mem"] = []
            gpu_series[f"gpu_{idx}_temp"] = []

        for m in metrics_list:
            gpus = {gpu["index"]: gpu for gpu in m.get("gpu_metrics", [])}
            for idx in gpu_indices:
                gpu = gpus.get(idx, {})
                gpu_series[f"gpu_{idx}_util"].append(gpu.get("utilization"))
                gpu_series[f"gpu_{idx}_mem"].append(gpu.get("memory_utilization"))
                gpu_series[f"gpu_{idx}_temp"].append(gpu.get("temperature"))

    return {
        "timestamps": [m["timestamp"] for m in metrics_list],
        "cpu_usage": [m["cpu_percent"] for m in metrics_list],
        "memory_usage": [m["memory_percent"] for m in metrics_list],
        "active_tasks": [m["active_tasks"] for m in metrics_list],
        "gpu_metrics": gpu_series
    }

@app.get("/api/gpu/info")
async def get_gpu_info():
    """Get current GPU information"""
    return {"gpus": get_gpu_metrics()}

# ---------------- Scripts listing ----------------
@app.get("/api/scripts/list")
async def scripts_list() -> Dict[str, Any]:
    root = WORKSPACE_DIR / "scripts"
    items: List[Dict[str, Any]] = []
    if root.exists():
        for p in sorted(root.rglob("*.py")):
            if p.is_file() and not p.name.startswith('.'):
                try:
                    items.append({
                        "path": str(p.relative_to(WORKSPACE_DIR)),
                        "size": p.stat().st_size
                    })
                except Exception:
                    continue
    return {"scripts": items}

@app.get("/api/data/tree")
async def get_data_tree():


                gpu = gpus.get(idx, {})
                gpu_series[f"gpu_{idx}_util"].append(gpu.get("utilization"))
                gpu_series[f"gpu_{idx}_mem"].append(gpu.get("memory_utilization"))
                gpu_series[f"gpu_{idx}_temp"].append(gpu.get("temperature"))

    return {
        "timestamps": [m["timestamp"] for m in metrics_list],
        "cpu_usage": [m["cpu_percent"] for m in metrics_list],
        "memory_usage": [m["memory_percent"] for m in metrics_list],
        "active_tasks": [m["active_tasks"] for m in metrics_list],
        "gpu_metrics": gpu_series
    }

@app.get("/api/gpu/info")
async def get_gpu_info():
    """Get current GPU information"""
    return {"gpus": get_gpu_metrics()}

# ---------------- Scripts listing ----------------
@app.get("/api/scripts/list")
async def scripts_list() -> Dict[str, Any]:
    root = WORKSPACE_DIR / "scripts"
    items: List[Dict[str, Any]] = []
    if root.exists():
        for p in sorted(root.rglob("*.py")):
            if p.is_file() and not p.name.startswith('.'):
                try:
                    items.append({
                        "path": str(p.relative_to(WORKSPACE_DIR)),
                        "size": p.stat().st_size
                    })
                except Exception:
                    continue
    return {"scripts": items}

@app.get("/api/data/tree")
async def get_data_tree():
    """Get data directory structure"""
    data_dir = WORKSPACE_DIR / "data"

    def build_tree(path: Path) -> Union[Dict, str]:
        if path.is_file():
            return str(path.relative_to(WORKSPACE_DIR))

        result = {}
        for item in path.iterdir():
            if item.name.startswith('.'):
                continue
            result[item.name] = build_tree(item)
        return result

    return build_tree(data_dir)

@app.get("/api/data/preview")