ApproveRequest

            except Exception as e:
                return {"text": f"plan error: {e}"}

        if cmd == "/verify":
            # Usage: /verify <json_trace>
            try:
                trace_json = _json_try(" ".join(args)) or {}
                vr = VerifyRequest(trace=trace_json)
                res = await api_verifier_verify(vr)
                return {"text": json.dumps(res, ensure_ascii=False)[:2000], "data": res}
            except Exception as e:
                return {"text": f"verify error: {e}"}

        return {"text": "Unknown command. Type /help."}

    # Fallback: treat plain text as inference
    out = await model_infer(InferRequest(prompt=text))
    return {"text": out.get("output", "")}

# ================= HITL Approval APIs =================
class ApproveRequest(BaseModel):
    id: str
    decision: str  # "approved" | "denied" | "modified" | "deferred"
    note: Optional[str] = None
    actor: Optional[str] = "api"

@app.get("/api/hitl/pending")
async def hitl_list_pending(kind: Optional[str] = None, sweep_ttl: Optional[int] = None, include_decided: bool = False) -> Dict[str, Any]:
    """List HITL items with optional kind filter. By default, shows only undecided pending items.

    Query params:
    - kind: filter by item kind (e.g., "curriculum", "code.patch")
    - sweep_ttl: if provided, sweep/expire pending items older than TTL seconds before listing (defaults from env or 600)
    - include_decided: if true, include items that already have a decision; else hide them
    """
    try:
        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        # Optional sweep to expire long-pending items
        try:
            ttl = int(sweep_ttl) if sweep_ttl is not None else int(os.environ.get("AGI_HITL_TTL", "600") or 600)
            if ttl > 0:
                _ = q.sweep_expired(ttl_sec=ttl, actor="api")
        except Exception:


                vr = VerifyRequest(trace=trace_json)
                res = await api_verifier_verify(vr)
                return {"text": json.dumps(res, ensure_ascii=False)[:2000], "data": res}
            except Exception as e:
                return {"text": f"verify error: {e}"}

        return {"text": "Unknown command. Type /help."}

    # Fallback: treat plain text as inference
    out = await model_infer(InferRequest(prompt=text))
    return {"text": out.get("output", "")}

# ================= HITL Approval APIs =================
class ApproveRequest(BaseModel):
    id: str
    decision: str  # "approved" | "denied" | "modified" | "deferred"
    note: Optional[str] = None
    actor: Optional[str] = "api"

@app.get("/api/hitl/pending")
async def hitl_list_pending(kind: Optional[str] = None, sweep_ttl: Optional[int] = None, include_decided: bool = False) -> Dict[str, Any]:
    """List HITL items with optional kind filter. By default, shows only undecided pending items.

    Query params:
    - kind: filter by item kind (e.g., "curriculum", "code.patch")
    - sweep_ttl: if provided, sweep/expire pending items older than TTL seconds before listing (defaults from env or 600)
    - include_decided: if true, include items that already have a decision; else hide them
    """
    try:
        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        # Optional sweep to expire long-pending items
        try:
            ttl = int(sweep_ttl) if sweep_ttl is not None else int(os.environ.get("AGI_HITL_TTL", "600") or 600)
            if ttl > 0:
                _ = q.sweep_expired(ttl_sec=ttl, actor="api")
        except Exception:
            pass
        items = q.read_pending()
        # Kind filter
        if kind:
            items = [it for it in items if str(it.get("kind", "")) == str(kind)]
        # Exclude those already decided unless requested
        if not include_decided:
            filtered = []
            for it in items:
                try:
                    if not q._has_decision(str(it.get("id", ""))):
                        filtered.append(it)
                except Exception:
                    filtered.append(it)
            items = filtered
        # Sort by timestamp desc when available
        try:
            items.sort(key=lambda x: str(x.get("ts", "")), reverse=True)
        except Exception:
            pass
        return {"pending": items}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/hitl/approve")
async def hitl_approve(req: ApproveRequest) -> Dict[str, Any]:
    """Approve/deny a HITL item by ID by writing to decisions.jsonl."""
    try:
        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ================= Selfplay APIs =================
class SelfplayStartRequest(BaseModel):
    # Optional overrides for env-driven knobs
    n_samples: Optional[int] = None
    max_tokens: Optional[int] = None
    lr: Optional[float] = None
    train_steps: Optional[int] = None


        # Exclude those already decided unless requested
        if not include_decided:
            filtered = []
            for it in items:
                try:
                    if not q._has_decision(str(it.get("id", ""))):
                        filtered.append(it)
                except Exception:
                    filtered.append(it)
            items = filtered
        # Sort by timestamp desc when available
        try:
            items.sort(key=lambda x: str(x.get("ts", "")), reverse=True)
        except Exception:
            pass
        return {"pending": items}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/hitl/approve")
async def hitl_approve(req: ApproveRequest) -> Dict[str, Any]:
    """Approve/deny a HITL item by ID by writing to decisions.jsonl."""
    try:
        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ================= Selfplay APIs =================
class SelfplayStartRequest(BaseModel):
    # Optional overrides for env-driven knobs
    n_samples: Optional[int] = None
    max_tokens: Optional[int] = None
    lr: Optional[float] = None
    train_steps: Optional[int] = None
    neg_steps: Optional[int] = None
    log_prompts: Optional[bool] = None
    mem_path: Optional[str] = None

@app.post("/api/selfplay/start")
async def selfplay_start(payload: SelfplayStartRequest, background_tasks: BackgroundTasks) -> Dict[str, Any]:
    """Start in-process selfplay loop as a background task via make target."""
    try:
        # Build environment overrides
        env = os.environ.copy()
        if payload.n_samples is not None:
            env["SELFPLAY_N_SAMPLES"] = str(int(payload.n_samples))


SelfplayStartRequest

            items.sort(key=lambda x: str(x.get("ts", "")), reverse=True)
        except Exception:
            pass
        return {"pending": items}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/hitl/approve")
async def hitl_approve(req: ApproveRequest) -> Dict[str, Any]:
    """Approve/deny a HITL item by ID by writing to decisions.jsonl."""
    try:
        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ================= Selfplay APIs =================
class SelfplayStartRequest(BaseModel):
    # Optional overrides for env-driven knobs
    n_samples: Optional[int] = None
    max_tokens: Optional[int] = None
    lr: Optional[float] = None
    train_steps: Optional[int] = None
    neg_steps: Optional[int] = None
    log_prompts: Optional[bool] = None
    mem_path: Optional[str] = None

@app.post("/api/selfplay/start")
async def selfplay_start(payload: SelfplayStartRequest, background_tasks: BackgroundTasks) -> Dict[str, Any]:
    """Start in-process selfplay loop as a background task via make target."""
    try:
        # Build environment overrides
        env = os.environ.copy()
        if payload.n_samples is not None:
            env["SELFPLAY_N_SAMPLES"] = str(int(payload.n_samples))
        if payload.max_tokens is not None:
            env["SELFPLAY_MAX_TOKENS"] = str(int(payload.max_tokens))
        if payload.lr is not None:
            env["SELFPLAY_LR"] = str(float(payload.lr))
        if payload.train_steps is not None:
            env["SELFPLAY_TRAIN_STEPS"] = str(int(payload.train_steps))
        if payload.neg_steps is not None:
            env["SELFPLAY_NEG_STEPS"] = str(int(payload.neg_steps))
        if payload.log_prompts is not None:
            env["SELFPLAY_LOG_PROMPTS"] = "1" if bool(payload.log_prompts) else "0"
        if payload.mem_path is not None and str(payload.mem_path).strip():


        root = WORKSPACE_DIR
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(root)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ================= Selfplay APIs =================
class SelfplayStartRequest(BaseModel):
    # Optional overrides for env-driven knobs
    n_samples: Optional[int] = None
    max_tokens: Optional[int] = None
    lr: Optional[float] = None
    train_steps: Optional[int] = None
    neg_steps: Optional[int] = None
    log_prompts: Optional[bool] = None
    mem_path: Optional[str] = None

@app.post("/api/selfplay/start")
async def selfplay_start(payload: SelfplayStartRequest, background_tasks: BackgroundTasks) -> Dict[str, Any]:
    """Start in-process selfplay loop as a background task via make target."""
    try:
        # Build environment overrides
        env = os.environ.copy()
        if payload.n_samples is not None:
            env["SELFPLAY_N_SAMPLES"] = str(int(payload.n_samples))
        if payload.max_tokens is not None:
            env["SELFPLAY_MAX_TOKENS"] = str(int(payload.max_tokens))
        if payload.lr is not None:
            env["SELFPLAY_LR"] = str(float(payload.lr))
        if payload.train_steps is not None:
            env["SELFPLAY_TRAIN_STEPS"] = str(int(payload.train_steps))
        if payload.neg_steps is not None:
            env["SELFPLAY_NEG_STEPS"] = str(int(payload.neg_steps))
        if payload.log_prompts is not None:
            env["SELFPLAY_LOG_PROMPTS"] = "1" if bool(payload.log_prompts) else "0"
        if payload.mem_path is not None and str(payload.mem_path).strip():
            env["SELFPLAY_MEM_PATH"] = str(payload.mem_path)
        else:
            # Ensure progress tails in UI by defaulting under data/
            env.setdefault("SELFPLAY_MEM_PATH", str(WORKSPACE_DIR / "data" / "selfplay" / "progress.jsonl"))
        # Launch make target in background
        cmd = ["make", "-C", str(WORKSPACE_DIR), "selfplay.run"]
        proc = await asyncio.create_subprocess_exec(*cmd, env=env)
        task_id = f"selfplay_{int(time.time())}"
        now = time.time()
        active_tasks[task_id] = {
            "kind": "selfplay",
            "pid": proc.pid,
            "started": now,
            "created_at": now,
            "status": "running",
            "cmd": " ".join(cmd),
        }
        task_processes[task_id] = proc
        return {"ok": True, "task_id": task_id, "pid": proc.pid}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/selfplay/stop")
async def selfplay_stop() -> Dict[str, Any]:
    """Stop running selfplay tasks."""
    try:
        stopped: List[str] = []
        for tid, proc in list(task_processes.items()):
            meta = active_tasks.get(tid) or {}
            if str(meta.get("kind")) != "selfplay":
                continue
            try:
                proc.terminate()
            except Exception:
                pass
            try:
                proc.kill()
            except Exception:
                pass
        # Mark status for downstream metrics
        try:


            env.setdefault("SELFPLAY_MEM_PATH", str(WORKSPACE_DIR / "data" / "selfplay" / "progress.jsonl"))
        # Launch make target in background
        cmd = ["make", "-C", str(WORKSPACE_DIR), "selfplay.run"]
        proc = await asyncio.create_subprocess_exec(*cmd, env=env)
        task_id = f"selfplay_{int(time.time())}"
        now = time.time()
        active_tasks[task_id] = {
            "kind": "selfplay",
            "pid": proc.pid,
            "started": now,
            "created_at": now,
            "status": "running",
            "cmd": " ".join(cmd),
        }
        task_processes[task_id] = proc
        return {"ok": True, "task_id": task_id, "pid": proc.pid}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/selfplay/stop")
async def selfplay_stop() -> Dict[str, Any]:
    """Stop running selfplay tasks."""
    try:
        stopped: List[str] = []
        for tid, proc in list(task_processes.items()):
            meta = active_tasks.get(tid) or {}
            if str(meta.get("kind")) != "selfplay":
                continue
            try:
                proc.terminate()
            except Exception:
                pass
            try:
                proc.kill()
            except Exception:
                pass
        # Mark status for downstream metrics
        try:
            meta["status"] = "stopped"
            meta["ended"] = time.time()
            active_tasks[tid] = meta
        except Exception:
            pass
            stopped.append(tid)
            task_processes.pop(tid, None)
            active_tasks.pop(tid, None)
        return {"ok": True, "stopped": stopped}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/selfplay/status")
async def selfplay_status() -> Dict[str, Any]:
    """Return current selfplay task status and recent GPU info."""
    items: List[Dict[str, Any]] = []
    for tid, meta in active_tasks.items():
        if str(meta.get("kind")) == "selfplay":
            state = {"task_id": tid, **meta}
            try:
                proc = task_processes.get(tid)
                if proc is not None and hasattr(proc, "returncode"):
                    state["returncode"] = proc.returncode
            except Exception:
                pass
            items.append(state)
    # Attach lightweight GPU info
    try:
        from agi_dw.core.utils.gpu_info import get_gpu_info  # type: ignore
        gpu = get_gpu_info()
    except Exception:


                pass
            try:
                proc.kill()
            except Exception:
                pass
        # Mark status for downstream metrics
        try:
            meta["status"] = "stopped"
            meta["ended"] = time.time()
            active_tasks[tid] = meta
        except Exception:
            pass
            stopped.append(tid)
            task_processes.pop(tid, None)
            active_tasks.pop(tid, None)
        return {"ok": True, "stopped": stopped}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/selfplay/status")
async def selfplay_status() -> Dict[str, Any]:
    """Return current selfplay task status and recent GPU info."""
    items: List[Dict[str, Any]] = []
    for tid, meta in active_tasks.items():
        if str(meta.get("kind")) == "selfplay":
            state = {"task_id": tid, **meta}
            try:
                proc = task_processes.get(tid)
                if proc is not None and hasattr(proc, "returncode"):
                    state["returncode"] = proc.returncode
            except Exception:
                pass
            items.append(state)
    # Attach lightweight GPU info
    try:
        from agi_dw.core.utils.gpu_info import get_gpu_info  # type: ignore
        gpu = get_gpu_info()
    except Exception:
        gpu = None
    return {"tasks": items, "gpu": gpu}

@app.get("/api/selfplay/progress")
async def selfplay_progress(limit: int = 200) -> Dict[str, Any]:
    """Tail progress JSONL file if present. Defaults to data/selfplay/progress.jsonl when available."""
    try:
        # Default to data/selfplay/progress.jsonl; fall back to root progress.jsonl
        cand = [WORKSPACE_DIR / "data" / "selfplay" / "progress.jsonl", WORKSPACE_DIR / "progress.jsonl"]
        path = None
        for p in cand:
            if p.exists():
                path = p
                break
        if path is None:
            return {"progress": []}
        rows: List[Dict[str, Any]] = []
        with path.open("r", encoding="utf-8") as f:
            lines = f.readlines()
        for li in lines[-max(1, int(limit)):]:
            li = (li or "").strip()
            if not li:


    items: List[Dict[str, Any]] = []
    for tid, meta in active_tasks.items():
        if str(meta.get("kind")) == "selfplay":
            state = {"task_id": tid, **meta}
            try:
                proc = task_processes.get(tid)
                if proc is not None and hasattr(proc, "returncode"):
                    state["returncode"] = proc.returncode
            except Exception:
                pass
            items.append(state)
    # Attach lightweight GPU info
    try:
        from agi_dw.core.utils.gpu_info import get_gpu_info  # type: ignore
        gpu = get_gpu_info()
    except Exception:
        gpu = None
    return {"tasks": items, "gpu": gpu}

@app.get("/api/selfplay/progress")
async def selfplay_progress(limit: int = 200) -> Dict[str, Any]:
    """Tail progress JSONL file if present. Defaults to data/selfplay/progress.jsonl when available."""
    try:
        # Default to data/selfplay/progress.jsonl; fall back to root progress.jsonl
        cand = [WORKSPACE_DIR / "data" / "selfplay" / "progress.jsonl", WORKSPACE_DIR / "progress.jsonl"]
        path = None
        for p in cand:
            if p.exists():
                path = p
                break
        if path is None:
            return {"progress": []}
        rows: List[Dict[str, Any]] = []
        with path.open("r", encoding="utf-8") as f:
            lines = f.readlines()
        for li in lines[-max(1, int(limit)):]:
            li = (li or "").strip()
            if not li:
                continue
            try:
                rows.append(json.loads(li))
            except Exception:
                continue
        return {"path": str(path.relative_to(WORKSPACE_DIR)), "progress": rows}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Selfplay HITL helpers (wrap existing HITL APIs with selfplay tag)
@app.get("/api/selfplay/hitl/pending")
async def selfplay_hitl_pending(sweep_ttl: Optional[int] = None) -> Dict[str, Any]:
    """List only selfplay curriculum items still pending a decision."""
    try:
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(WORKSPACE_DIR)
        # Optional sweep
        try:
            ttl = int(sweep_ttl) if sweep_ttl is not None else int(os.environ.get("AGI_HITL_TTL", "600") or 600)
            if ttl > 0:
                _ = q.sweep_expired(ttl_sec=ttl, actor="api")
        except Exception:
            pass
        items = q.read_pending()
        # Keep curriculum only and undecided only
        out = []
        for it in items:
            if str(it.get("kind", "")) != "curriculum":


                break
        if path is None:
            return {"progress": []}
        rows: List[Dict[str, Any]] = []
        with path.open("r", encoding="utf-8") as f:
            lines = f.readlines()
        for li in lines[-max(1, int(limit)):]:
            li = (li or "").strip()
            if not li:
                continue
            try:
                rows.append(json.loads(li))
            except Exception:
                continue
        return {"path": str(path.relative_to(WORKSPACE_DIR)), "progress": rows}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Selfplay HITL helpers (wrap existing HITL APIs with selfplay tag)
@app.get("/api/selfplay/hitl/pending")
async def selfplay_hitl_pending(sweep_ttl: Optional[int] = None) -> Dict[str, Any]:
    """List only selfplay curriculum items still pending a decision."""
    try:
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(WORKSPACE_DIR)
        # Optional sweep
        try:
            ttl = int(sweep_ttl) if sweep_ttl is not None else int(os.environ.get("AGI_HITL_TTL", "600") or 600)
            if ttl > 0:
                _ = q.sweep_expired(ttl_sec=ttl, actor="api")
        except Exception:
            pass
        items = q.read_pending()
        # Keep curriculum only and undecided only
        out = []
        for it in items:
            if str(it.get("kind", "")) != "curriculum":
                continue
            try:
                if q._has_decision(str(it.get("id", ""))):
                    continue
            except Exception:
                pass
            out.append(it)
        try:
            out.sort(key=lambda x: str(x.get("ts", "")), reverse=True)
        except Exception:
            pass
        return {"pending": out}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/selfplay/hitl/approve")
async def selfplay_hitl_approve(req: ApproveRequest) -> Dict[str, Any]:
    try:
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(WORKSPACE_DIR)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


        # Keep curriculum only and undecided only
        out = []
        for it in items:
            if str(it.get("kind", "")) != "curriculum":
                continue
            try:
                if q._has_decision(str(it.get("id", ""))):
                    continue
            except Exception:
                pass
            out.append(it)
        try:
            out.sort(key=lambda x: str(x.get("ts", "")), reverse=True)
        except Exception:
            pass
        return {"pending": out}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/selfplay/hitl/approve")
async def selfplay_hitl_approve(req: ApproveRequest) -> Dict[str, Any]:
    try:
        from agi_dw.core.hitl.approval_queue import ApprovalQueue  # type: ignore
        q = ApprovalQueue(WORKSPACE_DIR)
        dec = q.update_status(item_id=str(req.id), decision=str(req.decision), note=(req.note or None), actor=(req.actor or "api"))
        return {"ok": True, "decision": dec}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


                    '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


    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:


generate

# 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):


chat

        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(),


dispatch

            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({


# 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)


allocate

        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


release

            # 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,


        """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",


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")
async def preview_file(path: str):
    """Preview file contents"""
    full_path = WORKSPACE_DIR / path

    # Basic security check
    if not str(full_path).startswith(str(WORKSPACE_DIR)):
        raise HTTPException(status_code=403, detail="Access denied")

    try:
        with open(full_path, 'r') as f:
            # Read first 100KB for preview
            content = f.read(102400)
            return {"content": content}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


    model_cache.clear()

    return {"status": "cleanup completed"}

# Secure file download (allowlist: under data/ and api_server.log)
@app.get("/api/file/download")
async def file_download(path: str) -> FileResponse:
    full = (WORKSPACE_DIR / path).resolve()
    # allow only data/* or api log
    if not (str(full).startswith(str((WORKSPACE_DIR / "data").resolve())) or full.name == "api_server.log"):
        raise HTTPException(status_code=403, detail="Download not allowed")
    if not full.exists() or not full.is_file():
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(str(full), filename=full.name)

# Benchmark comparison endpoint
@app.get("/api/bench/solution_diff")
async def get_solution_diff(suite: str, run_a: str, run_b: str, task_id: str) -> Dict[str, Any]:
    """Get solution diff for a specific task between two runs"""
    # Load run metas
    def _load_solutions(rd: str) -> Dict[str, str]:
        solutions = {}
        rdir = _runs_dir_for_suite(suite) / rd
        rj = rdir / "run.json"
        if rj.exists():
            meta = _read_json_file(rj)
            paths = meta.get("paths", {}) if isinstance(meta, dict) else {}
            results_path = paths.get("results")
            if isinstance(results_path, str) and results_path:
                p = WORKSPACE_DIR / results_path
                if p.exists():
                    try:
                        with p.open("r") as f:
                            for line in f:
                                try:
                                    obj = json.loads(line.strip())
                                    if obj.get("task_id") == task_id:
                                        solutions[task_id] = obj.get("solution", "")
                                except Exception:
                                    continue
                    except Exception:
                        pass
        return solutions

    solutions_a = _load_solutions(run_a)
    solutions_b = _load_solutions(run_b)

    return {
        "task_id": task_id,
        "solution_a": solutions_a.get(task_id),
        "solution_b": solutions_b.get(task_id)
    }

@app.get("/api/bench/compare")
async def bench_compare(suite: str, run_a: str, run_b: str) -> Dict[str, Any]:
    # Load run metas
    def _load_meta(rd: str) -> Dict[str, Any]:
        p = _runs_dir_for_suite(suite) / rd / "run.json"
        return _read_json_file(p) if p.exists() else {}
    a_meta = _load_meta(run_a)
    b_meta = _load_meta(run_b)
    # Load per-task flags using same logic as bench_run_tasks
    async def _load_tasks(rd: str) -> Dict[str, bool]:


                                    if obj.get("task_id") == task_id:
                                        solutions[task_id] = obj.get("solution", "")
                                except Exception:
                                    continue
                    except Exception:
                        pass
        return solutions

    solutions_a = _load_solutions(run_a)
    solutions_b = _load_solutions(run_b)

    return {
        "task_id": task_id,
        "solution_a": solutions_a.get(task_id),
        "solution_b": solutions_b.get(task_id)
    }

@app.get("/api/bench/compare")
async def bench_compare(suite: str, run_a: str, run_b: str) -> Dict[str, Any]:
    # Load run metas
    def _load_meta(rd: str) -> Dict[str, Any]:
        p = _runs_dir_for_suite(suite) / rd / "run.json"
        return _read_json_file(p) if p.exists() else {}
    a_meta = _load_meta(run_a)
    b_meta = _load_meta(run_b)
    # Load per-task flags using same logic as bench_run_tasks
    async def _load_tasks(rd: str) -> Dict[str, bool]:
        rdir = _runs_dir_for_suite(suite) / rd
        rj = rdir / "run.json"
        passed: Dict[str, bool] = {}
        if rj.exists():
            meta = _read_json_file(rj)
            paths = meta.get("paths", {}) if isinstance(meta, dict) else {}
            out_path = paths.get("out")
            results_path = paths.get("results")
            if isinstance(out_path, str) and out_path:
                p = WORKSPACE_DIR / out_path
                if p.exists():
                    try:
                        with p.open("r", encoding="utf-8") as f:
                            for line in f:
                                try:
                                    obj = json.loads(line.strip())


        return solutions

    solutions_a = _load_solutions(run_a)
    solutions_b = _load_solutions(run_b)

    return {
        "task_id": task_id,
        "solution_a": solutions_a.get(task_id),
        "solution_b": solutions_b.get(task_id)
    }

@app.get("/api/bench/compare")
async def bench_compare(suite: str, run_a: str, run_b: str) -> Dict[str, Any]:
    # Load run metas
    def _load_meta(rd: str) -> Dict[str, Any]:
        p = _runs_dir_for_suite(suite) / rd / "run.json"
        return _read_json_file(p) if p.exists() else {}
    a_meta = _load_meta(run_a)
    b_meta = _load_meta(run_b)
    # Load per-task flags using same logic as bench_run_tasks
    async def _load_tasks(rd: str) -> Dict[str, bool]:
        rdir = _runs_dir_for_suite(suite) / rd
        rj = rdir / "run.json"
        passed: Dict[str, bool] = {}
        if rj.exists():
            meta = _read_json_file(rj)
            paths = meta.get("paths", {}) if isinstance(meta, dict) else {}
            out_path = paths.get("out")
            results_path = paths.get("results")
            if isinstance(out_path, str) and out_path:
                p = WORKSPACE_DIR / out_path
                if p.exists():
                    try:
                        with p.open("r", encoding="utf-8") as f:
                            for line in f:
                                try:
                                    obj = json.loads(line.strip())
                                    if obj.get("suite") == suite:
                                        passed[str(obj.get("task_id"))] = bool(obj.get("pass1", False))
                                except Exception:
                                    continue
                    except Exception:
                        pass
            if not passed and isinstance(results_path, str) and results_path:
                p = WORKSPACE_DIR / results_path
                if p.exists():
                    try:
                        with p.open("r", encoding="utf-8") as f:
                            for line in f:
                                try:
                                    obj = json.loads(line.strip())
                                    passed[str(obj.get("task_id"))] = bool(obj.get("passed", False))
                                except Exception:
                                    continue
                    except Exception:
                        pass
        return passed
    a_tasks = await _load_tasks(run_a)
    b_tasks = await _load_tasks(run_b)
    all_tids = sorted(set(a_tasks.keys()) | set(b_tasks.keys()))
    diffs = []
    changed = 0
    for tid in all_tids:
        a = bool(a_tasks.get(tid, False))
        b = bool(b_tasks.get(tid, False))
        if a != b:
            changed += 1
        diffs.append({"task_id": tid, "run_a": a, "run_b": b, "delta": (1 if b and not a else (-1 if a and not b else 0))})
    # Aggregate deltas
    def _get_pass1(meta: Dict[str, Any]) -> Optional[float]:
        m = meta.get("metrics", {}) if isinstance(meta, dict) else {}
        v = m.get("pass1_rate") if isinstance(m, dict) else None
        try:
            return float(v) if v is not None else None
        except Exception:
            return None
    pass1_a = _get_pass1(a_meta)


                                try:
                                    obj = json.loads(line.strip())
                                    passed[str(obj.get("task_id"))] = bool(obj.get("passed", False))
                                except Exception:
                                    continue
                    except Exception:
                        pass
        return passed
    a_tasks = await _load_tasks(run_a)
    b_tasks = await _load_tasks(run_b)
    all_tids = sorted(set(a_tasks.keys()) | set(b_tasks.keys()))
    diffs = []
    changed = 0
    for tid in all_tids:
        a = bool(a_tasks.get(tid, False))
        b = bool(b_tasks.get(tid, False))
        if a != b:
            changed += 1
        diffs.append({"task_id": tid, "run_a": a, "run_b": b, "delta": (1 if b and not a else (-1 if a and not b else 0))})
    # Aggregate deltas
    def _get_pass1(meta: Dict[str, Any]) -> Optional[float]:
        m = meta.get("metrics", {}) if isinstance(meta, dict) else {}
        v = m.get("pass1_rate") if isinstance(m, dict) else None
        try:
            return float(v) if v is not None else None
        except Exception:
            return None
    pass1_a = _get_pass1(a_meta)
    pass1_b = _get_pass1(b_meta)
    delta_pass1 = None
    if pass1_a is not None and pass1_b is not None:
        delta_pass1 = pass1_b - pass1_a
    return {
        "suite": suite,
        "run_a": {"dir": run_a, "pass1_rate": pass1_a},
        "run_b": {"dir": run_b, "pass1_rate": pass1_b},
        "delta_pass1": delta_pass1,
        "changed": changed,
        "diffs": diffs,
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)

# ================= Model Registry APIs =================


                pass

        # Render Markdown to HTML
        html = markdown.markdown(content, extensions=['tables', 'fenced_code'])

        return {
            "path": str(full.relative_to(WORKSPACE_DIR)),
            "meta": meta,
            "html": html,
            "raw": content
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/task/logs/stream/{task_id}")
async def stream_task_logs(task_id: str) -> StreamingResponse:
    """Stream task logs as Server-Sent Events."""
    if task_id not in active_tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    async def event_generator() -> AsyncGenerator[str, None]:
        last_stdout = ""
        last_stderr = ""
        while True:
            task = active_tasks.get(task_id)
            if not task:
                break

            stdout = task.get("stdout", "")
            stderr = task.get("stderr", "")

            if stdout != last_stdout:
                yield f"data: {json.dumps({'type': 'stdout', 'data': stdout[len(last_stdout):]})}\n\n"
                last_stdout = stdout

            if stderr != last_stderr:
                yield f"data: {json.dumps({'type': 'stderr', 'data': stderr[len(last_stderr):]})}\n\n"
                last_stderr = stderr

            if task.get("status") not in ("running", "queued"):
                break

            await asyncio.sleep(1.0)

    return StreamingResponse(event_generator(), media_type="text/event-stream")

# ---------------- Dev endpoints (patch/test) ----------------
class PatchRequest(BaseModel):
    repo: str
    diff: str

class TestRequest(BaseModel):
    repo: str
    pytest_args: Optional[str] = "-q"

@app.post("/api/dev/patch/dry_run")
async def dev_patch_dry(req: PatchRequest) -> Dict[str, Any]:
    try:
        from agi_dw.tools.patch_policy import load_env_limits, validate_unified_diff  # type: ignore
        ok, why, meta = validate_unified_diff(req.diff, load_env_limits(strict_default=True))
        policy = {"ok": bool(ok), "why": why, **(meta or {})}
    except Exception:
        policy = {"ok": True}


@app.post("/api/dev/compose")
async def dev_compose(req: ComposeRequest) -> Dict[str, Any]:
    root = WORKSPACE_DIR
    trace_dir = root / "data" / "traces" / "compose"
    trace_dir.mkdir(parents=True, exist_ok=True)
    # Idempotency: reuse run id if key provided and a file exists
    if req.idempotency_key:
        safe_key = re.sub(r"[^a-zA-Z0-9_-]", "_", req.idempotency_key)
        run_id = safe_key
    else:
        run_id = datetime.now().strftime("%Y%m%dT%H%M%SZ") + "_" + uuid.uuid4().hex[:6]
    out = trace_dir / f"{run_id}.jsonl"
    if req.idempotency_key and out.exists():
        # Return existing summary tail
        try:
            with out.open("r", encoding="utf-8") as f:
                lines = f.readlines()[-50:]
            return {"run_id": run_id, "trace": str(out.relative_to(root)), "status": "idempotent_reuse", "tail": [l.strip() for l in lines]}
        except Exception:
            return {"run_id": run_id, "trace": str(out.relative_to(root)), "status": "idempotent_reuse"}
    def _write(obj: Dict[str, Any]) -> None:
        try:
            with out.open("a", encoding="utf-8") as f:
                f.write(json.dumps(obj, ensure_ascii=False) + "\n")
        except Exception:
            pass
    summary: List[Dict[str, Any]] = []
    for i, st in enumerate(req.steps):
        rec_base = {"i": i, "tool": st.tool, "args": st.args, "ts": datetime.now().isoformat()}
        try:
            if st.tool == "run_cmd":
                argv = st.args.get("argv") or []
                cwd = st.args.get("cwd") or str(root)
                p = await asyncio.create_subprocess_exec(*argv, cwd=str(cwd), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
                out_b, err_b = await p.communicate()
                res = {"rc": p.returncode, "stdout": out_b.decode(errors="replace"), "stderr": err_b.decode(errors="replace")}
                _write({**rec_base, "result": res})
                summary.append({"tool": st.tool, "ok": p.returncode == 0})
            elif st.tool == "apply_patch":
                from agi_dw.tools.patch_actuator import apply_unified_diff  # type: ignore
                res = apply_unified_diff(st.args.get("repo"), st.args.get("diff", ""), st.args.get("allow_globs", ["**/*"]), st.args.get("block_globs", ["**/.ssh/**"]), max_files=int(st.args.get("max_files", 50)), dry_run=bool(st.args.get("dry_run", False)))
                _write({**rec_base, "result": res})
                summary.append({"tool": st.tool, "ok": bool(res.get("ok"))})
            elif st.tool == "dev_run_tests":
                res = await dev_run_tests(TestRequest(repo=str(st.args.get("repo")), pytest_args=str(st.args.get("pytest_args", "-q"))))
                _write({**rec_base, "result": res})


class ChatRequest(BaseModel):
    message: str
    session_id: Optional[str] = None

def _json_try(s: str) -> Optional[Dict[str, Any]]:
    try:
        return json.loads(s)
    except Exception:
        return None

@app.post("/api/agent/chat")
async def agent_chat(req: ChatRequest, background_tasks: BackgroundTasks, database: Session = Depends(db.get_db)) -> Dict[str, Any]:
    """
    Minimal AGI chat endpoint with slash-commands that invoke existing capabilities.
    Fallback to simple model inference.
    """
    text = (req.message or "").strip()
    if not text:
        return {"text": "Say something, or type /help for commands."}

    async def _ok(resp: Any) -> Dict[str, Any]:
        return {"text": "Done.", "data": resp}

    if text.startswith("/"):
        parts = text.split()
        cmd = parts[0].lower()
        args = parts[1:]

        if cmd in {"/help", "/?"}:
            return {"text": "Commands:\n"
                            "/make <target> [args_json]\n"
                            "/logs <task_id>\n"
                            "/bench runs <suite>|/bench compare <suite> <run_a> <run_b>\n"
                            "/qa <question>\n"
                            "/search <pattern> [base] [path]\n"
                            "/repo tree [path]\n"
                            "/file read <path> [base]\n"
                            "/models|/model load <id>|/infer <prompt>"}

        if cmd == "/make" and args:
            target = args[0]
            extra = " ".join(args[1:]).strip()


    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


        # 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,


async def api_verifier_verify(req: VerifyRequest) -> Dict[str, Any]:
    cfg = VerifierServiceConfig(
        model=str(req.model),
        backend=str(req.backend),
        adapter_dir=req.adapter_dir,
        adapter_bank=req.adapter_bank,
        structured_mode=str(req.structured_mode),
        timeout_sec=int(req.timeout_sec),
        strict=bool(req.strict),
        calibrate=bool(req.calibrate),
        calib_model=req.calib_model,
        log_prompts=bool(req.log_prompts),
    )
    return verifier_run(req.trace, cfg)

@app.post("/api/updater/run")
async def api_updater_run(req: UpdaterRequest) -> Dict[str, Any]:
    try:
        upd = Updater(repo_root=WORKSPACE_DIR, fast=bool(req.fast))
        # Run in background thread to avoid blocking
        def _run() -> None:
            try:
                upd.run()
            except Exception:
                pass
        threading.Thread(target=_run, daemon=True).start()
        return {"ok": True, "started": True}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ================= Additional central management APIs =================

# File system tree (repo/data)
def _safe_root(base: str) -> Path:
    if base == "data":
        return WORKSPACE_DIR / "data"
    return WORKSPACE_DIR

def _build_tree(path: Path, rel_to: Path, depth: int = 0, max_depth: int = 4, max_entries: int = 200) -> Union[Dict[str, Any], str]:
    if depth > max_depth:
        return "…"
    if path.is_file():
        return str(path.relative_to(rel_to))
    entries: Dict[str, Any] = {}
    try:


                            "/file read <path> [base]\n"
                            "/models|/model load <id>|/infer <prompt>"}

        if cmd == "/make" and args:
            target = args[0]
            extra = " ".join(args[1:]).strip()
            arg_obj = _json_try(extra) if extra else None
            req_obj = TaskRequest(task=target, args=arg_obj or None)
            res = await start_task(req_obj, background_tasks, database)
            return {"text": f"Started {target}", "data": res}

        if cmd == "/logs" and args:
            tid = args[0]
            try:
                logs = await get_task_logs(tid)  # type: ignore[arg-type]
            except HTTPException as e:
                return {"text": f"Error: {e.detail}"}
            # Trim for chat
            out = (logs.get("stdout", "") or "")
            err = (logs.get("stderr", "") or "")
            def _tail(s: str, n: int = 80) -> str:
                lines = s.splitlines()
                return "\n".join(lines[-n:])
            return {"text": f"STDOUT:\n{_tail(out)}\n\nSTDERR:\n{_tail(err)}"}

        if cmd == "/bench" and args:
            sub = args[0]
            if sub == "runs" and len(args) >= 2:
                suite = args[1]
                runs = await list_bench_runs(suite)
                return {"text": f"Found {len(runs)} runs.", "data": runs[:50]}
            if sub == "compare" and len(args) >= 4:
                suite, a, b = args[1], args[2], args[3]
                cmp = await bench_compare(suite, a, b)
                return await _ok(cmp)
            return {"text": "Usage: /bench runs <suite> | /bench compare <suite> <run_a> <run_b>"}

        if cmd == "/qa" and args:
            question = text[len("/qa"):].strip()
            ans = await qa_query({"question": question})
            return {"text": ans.get("answer", ""), "data": ans.get("citations")}

        if cmd == "/search" and args:


from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, JSON, ForeignKey, Enum as SQLEnum, or_, desc, asc, case
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from typing import Tuple
from datetime import datetime
from typing import Optional, Dict, Any, List
import enum
import json
import os

# Create SQLite database engine
DATABASE_URL = "sqlite:///./dashboard.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})

# Create declarative base
Base = declarative_base()

class TaskStatus(str, enum.Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"

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

class Task(Base):
    __tablename__ = "tasks"

    id = Column(String, primary_key=True)
    name = Column(String, nullable=False)
    status = Column(SQLEnum(TaskStatus), nullable=False)
    priority = Column(SQLEnum(TaskPriority), nullable=False, default=TaskPriority.NORMAL)
    args = Column(JSON, nullable=True)
    dependencies = Column(JSON, nullable=True)  # List of task IDs
    resources = Column(JSON, nullable=True)  # Resource requirements
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    return_code = Column(Integer, nullable=True)
    error = Column(String, nullable=True)
    stdout = Column(String, nullable=True)
    stderr = Column(String, nullable=True)
    progress = Column(Float, nullable=False, default=0)

    metrics = relationship("TaskMetric", back_populates="task", cascade="all, delete-orphan")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "status": self.status.value,
            "priority": self.priority.value,
            "args": self.args,
            "dependencies": self.dependencies or [],
            "resources": self.resources,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "return_code": self.return_code,
            "error": self.error,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "progress": self.progress
        }

class TaskMetric(Base):
    __tablename__ = "task_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(String, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    name = Column(String, nullable=False)
    value = Column(Float, nullable=False)
    meta_data = Column(JSON, nullable=True)  # Changed from metadata to meta_data to avoid SQLAlchemy reserved word

    task = relationship("Task", back_populates="metrics")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "task_id": self.task_id,
            "timestamp": self.timestamp.isoformat(),
            "name": self.name,
            "value": self.value,
            "metadata": self.meta_data  # Using metadata in API for consistency, but meta_data in DB
        }

class SystemMetric(Base):
    __tablename__ = "system_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    cpu_percent = Column(Float, nullable=True)
    memory_percent = Column(Float, nullable=True)
    active_tasks = Column(Integer, nullable=True)
    gpu_metrics = Column(JSON, nullable=True)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "timestamp": self.timestamp.isoformat(),
            "cpu_percent": self.cpu_percent,
            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }

# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()

def get_tasks(
    db: SessionLocal,
    status: Optional[str] = None,
    priority: Optional[str] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100
) -> Tuple[List[Task], int]:
    """Get tasks with filters, sorting and pagination"""
    query = db.query(Task)

    # Apply filters
    if status:
        query = query.filter(Task.status == status)
    if priority:
        query = query.filter(Task.priority == priority)
    if search:
        search = f"%{search}%"
        query = query.filter(
            or_(
                Task.name.ilike(search),
                Task.id.ilike(search),
                Task.error.ilike(search) if Task.error is not None else False
            )
        )

    # Get total count for pagination
    total = query.count()

    # Apply sorting
    if sort_by == "priority":
        # Custom priority ordering
        priority_case = case(
            [(Task.priority == TaskPriority.CRITICAL, 0),
             (Task.priority == TaskPriority.HIGH, 1),
             (Task.priority == TaskPriority.NORMAL, 2)],
            else_=3
        )
        query = query.order_by(priority_case if not sort_desc else priority_case.desc())
    else:
        # Default to created_at
        query = query.order_by(
            desc(Task.created_at) if sort_desc else asc(Task.created_at)
        )

    # Apply pagination
    query = query.offset(offset).limit(limit)

    return query.all(), total

def create_task(db: SessionLocal, task_data: Dict[str, Any]) -> Task:
    """Create a new task"""
    task = Task(**task_data)
    db.add(task)
    db.commit()
    db.refresh(task)
    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:
    """Add system metrics"""
    metrics = SystemMetric(**metrics_data)
    db.add(metrics)
    db.commit()
    db.refresh(metrics)
    return metrics

def get_system_metrics(
    db: SessionLocal,
    start_time: Optional[datetime] = None,
    end_time: Optional[datetime] = None,
    limit: int = 100
) -> List[SystemMetric]:
    """Get system metrics with optional time range"""
    query = db.query(SystemMetric)
    if start_time:
        query = query.filter(SystemMetric.timestamp >= start_time)
    if end_time:
        query = query.filter(SystemMetric.timestamp <= end_time)
    return query.order_by(SystemMetric.timestamp.desc()).limit(limit).all()


TaskStatus

from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, JSON, ForeignKey, Enum as SQLEnum, or_, desc, asc, case
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from typing import Tuple
from datetime import datetime
from typing import Optional, Dict, Any, List
import enum
import json
import os

# Create SQLite database engine
DATABASE_URL = "sqlite:///./dashboard.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})

# Create declarative base
Base = declarative_base()

class TaskStatus(str, enum.Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"

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

class Task(Base):
    __tablename__ = "tasks"

    id = Column(String, primary_key=True)
    name = Column(String, nullable=False)
    status = Column(SQLEnum(TaskStatus), nullable=False)
    priority = Column(SQLEnum(TaskPriority), nullable=False, default=TaskPriority.NORMAL)
    args = Column(JSON, nullable=True)
    dependencies = Column(JSON, nullable=True)  # List of task IDs
    resources = Column(JSON, nullable=True)  # Resource requirements
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)


TaskPriority

from datetime import datetime
from typing import Optional, Dict, Any, List
import enum
import json
import os

# Create SQLite database engine
DATABASE_URL = "sqlite:///./dashboard.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})

# Create declarative base
Base = declarative_base()

class TaskStatus(str, enum.Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"

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

class Task(Base):
    __tablename__ = "tasks"

    id = Column(String, primary_key=True)
    name = Column(String, nullable=False)
    status = Column(SQLEnum(TaskStatus), nullable=False)
    priority = Column(SQLEnum(TaskPriority), nullable=False, default=TaskPriority.NORMAL)
    args = Column(JSON, nullable=True)
    dependencies = Column(JSON, nullable=True)  # List of task IDs
    resources = Column(JSON, nullable=True)  # Resource requirements
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    return_code = Column(Integer, nullable=True)
    error = Column(String, nullable=True)
    stdout = Column(String, nullable=True)
    stderr = Column(String, nullable=True)
    progress = Column(Float, nullable=False, default=0)


Task

# Create SQLite database engine
DATABASE_URL = "sqlite:///./dashboard.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})

# Create declarative base
Base = declarative_base()

class TaskStatus(str, enum.Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"

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

class Task(Base):
    __tablename__ = "tasks"

    id = Column(String, primary_key=True)
    name = Column(String, nullable=False)
    status = Column(SQLEnum(TaskStatus), nullable=False)
    priority = Column(SQLEnum(TaskPriority), nullable=False, default=TaskPriority.NORMAL)
    args = Column(JSON, nullable=True)
    dependencies = Column(JSON, nullable=True)  # List of task IDs
    resources = Column(JSON, nullable=True)  # Resource requirements
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    return_code = Column(Integer, nullable=True)
    error = Column(String, nullable=True)
    stdout = Column(String, nullable=True)
    stderr = Column(String, nullable=True)
    progress = Column(Float, nullable=False, default=0)

    metrics = relationship("TaskMetric", back_populates="task", cascade="all, delete-orphan")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "status": self.status.value,
            "priority": self.priority.value,
            "args": self.args,
            "dependencies": self.dependencies or [],
            "resources": self.resources,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "return_code": self.return_code,
            "error": self.error,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "progress": self.progress
        }

class TaskMetric(Base):
    __tablename__ = "task_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(String, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    name = Column(String, nullable=False)
    value = Column(Float, nullable=False)
    meta_data = Column(JSON, nullable=True)  # Changed from metadata to meta_data to avoid SQLAlchemy reserved word

    task = relationship("Task", back_populates="metrics")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "task_id": self.task_id,
            "timestamp": self.timestamp.isoformat(),
            "name": self.name,
            "value": self.value,


TaskMetric

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "status": self.status.value,
            "priority": self.priority.value,
            "args": self.args,
            "dependencies": self.dependencies or [],
            "resources": self.resources,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "return_code": self.return_code,
            "error": self.error,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "progress": self.progress
        }

class TaskMetric(Base):
    __tablename__ = "task_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(String, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    name = Column(String, nullable=False)
    value = Column(Float, nullable=False)
    meta_data = Column(JSON, nullable=True)  # Changed from metadata to meta_data to avoid SQLAlchemy reserved word

    task = relationship("Task", back_populates="metrics")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "task_id": self.task_id,
            "timestamp": self.timestamp.isoformat(),
            "name": self.name,
            "value": self.value,
            "metadata": self.meta_data  # Using metadata in API for consistency, but meta_data in DB
        }

class SystemMetric(Base):
    __tablename__ = "system_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    cpu_percent = Column(Float, nullable=True)
    memory_percent = Column(Float, nullable=True)
    active_tasks = Column(Integer, nullable=True)
    gpu_metrics = Column(JSON, nullable=True)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "timestamp": self.timestamp.isoformat(),
            "cpu_percent": self.cpu_percent,
            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }


SystemMetric

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(String, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    name = Column(String, nullable=False)
    value = Column(Float, nullable=False)
    meta_data = Column(JSON, nullable=True)  # Changed from metadata to meta_data to avoid SQLAlchemy reserved word

    task = relationship("Task", back_populates="metrics")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "task_id": self.task_id,
            "timestamp": self.timestamp.isoformat(),
            "name": self.name,
            "value": self.value,
            "metadata": self.meta_data  # Using metadata in API for consistency, but meta_data in DB
        }

class SystemMetric(Base):
    __tablename__ = "system_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    cpu_percent = Column(Float, nullable=True)
    memory_percent = Column(Float, nullable=True)
    active_tasks = Column(Integer, nullable=True)
    gpu_metrics = Column(JSON, nullable=True)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "timestamp": self.timestamp.isoformat(),
            "cpu_percent": self.cpu_percent,
            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }

# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()


    id = Column(Integer, primary_key=True, autoincrement=True)
    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
    cpu_percent = Column(Float, nullable=True)
    memory_percent = Column(Float, nullable=True)
    active_tasks = Column(Integer, nullable=True)
    gpu_metrics = Column(JSON, nullable=True)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "timestamp": self.timestamp.isoformat(),
            "cpu_percent": self.cpu_percent,
            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }

# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()

def get_tasks(
    db: SessionLocal,
    status: Optional[str] = None,
    priority: Optional[str] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100
) -> Tuple[List[Task], int]:
    """Get tasks with filters, sorting and pagination"""


        return {
            "id": self.id,
            "timestamp": self.timestamp.isoformat(),
            "cpu_percent": self.cpu_percent,
            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }

# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()

def get_tasks(
    db: SessionLocal,
    status: Optional[str] = None,
    priority: Optional[str] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100
) -> Tuple[List[Task], int]:
    """Get tasks with filters, sorting and pagination"""
    query = db.query(Task)

    # Apply filters
    if status:


            "memory_percent": self.memory_percent,
            "active_tasks": self.active_tasks,
            "gpu_metrics": self.gpu_metrics
        }

# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()

def get_tasks(
    db: SessionLocal,
    status: Optional[str] = None,
    priority: Optional[str] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100
) -> Tuple[List[Task], int]:
    """Get tasks with filters, sorting and pagination"""
    query = db.query(Task)

    # Apply filters
    if status:
        query = query.filter(Task.status == status)
    if priority:
        query = query.filter(Task.priority == priority)
    if search:


# Create database session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """Get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def init_db():
    """Initialize database"""
    Base.metadata.create_all(bind=engine)

def get_task_by_id(db: SessionLocal, task_id: str) -> Optional[Task]:
    """Get task by ID"""
    return db.query(Task).filter(Task.id == task_id).first()

def get_tasks(
    db: SessionLocal,
    status: Optional[str] = None,
    priority: Optional[str] = None,
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_desc: bool = True,
    offset: int = 0,
    limit: int = 100
) -> Tuple[List[Task], int]:
    """Get tasks with filters, sorting and pagination"""
    query = db.query(Task)

    # Apply filters
    if status:
        query = query.filter(Task.status == status)
    if priority:
        query = query.filter(Task.priority == priority)
    if search:
        search = f"%{search}%"
        query = query.filter(
            or_(
                Task.name.ilike(search),
                Task.id.ilike(search),
                Task.error.ilike(search) if Task.error is not None else False
            )
        )

    # Get total count for pagination
    total = query.count()

    # Apply sorting
    if sort_by == "priority":
        # Custom priority ordering
        priority_case = case(
            [(Task.priority == TaskPriority.CRITICAL, 0),
             (Task.priority == TaskPriority.HIGH, 1),
             (Task.priority == TaskPriority.NORMAL, 2)],
            else_=3
        )
        query = query.order_by(priority_case if not sort_desc else priority_case.desc())
    else:
        # Default to created_at
        query = query.order_by(
            desc(Task.created_at) if sort_desc else asc(Task.created_at)
        )

    # Apply pagination
    query = query.offset(offset).limit(limit)

    return query.all(), total

def create_task(db: SessionLocal, task_data: Dict[str, Any]) -> Task:
    """Create a new task"""
    task = Task(**task_data)
    db.add(task)
    db.commit()
    db.refresh(task)
    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:


    if sort_by == "priority":
        # Custom priority ordering
        priority_case = case(
            [(Task.priority == TaskPriority.CRITICAL, 0),
             (Task.priority == TaskPriority.HIGH, 1),
             (Task.priority == TaskPriority.NORMAL, 2)],
            else_=3
        )
        query = query.order_by(priority_case if not sort_desc else priority_case.desc())
    else:
        # Default to created_at
        query = query.order_by(
            desc(Task.created_at) if sort_desc else asc(Task.created_at)
        )

    # Apply pagination
    query = query.offset(offset).limit(limit)

    return query.all(), total

def create_task(db: SessionLocal, task_data: Dict[str, Any]) -> Task:
    """Create a new task"""
    task = Task(**task_data)
    db.add(task)
    db.commit()
    db.refresh(task)
    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:
    """Add system metrics"""
    metrics = SystemMetric(**metrics_data)
    db.add(metrics)
    db.commit()
    db.refresh(metrics)
    return metrics

def get_system_metrics(


        query = query.order_by(priority_case if not sort_desc else priority_case.desc())
    else:
        # Default to created_at
        query = query.order_by(
            desc(Task.created_at) if sort_desc else asc(Task.created_at)
        )

    # Apply pagination
    query = query.offset(offset).limit(limit)

    return query.all(), total

def create_task(db: SessionLocal, task_data: Dict[str, Any]) -> Task:
    """Create a new task"""
    task = Task(**task_data)
    db.add(task)
    db.commit()
    db.refresh(task)
    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:
    """Add system metrics"""
    metrics = SystemMetric(**metrics_data)
    db.add(metrics)
    db.commit()
    db.refresh(metrics)
    return metrics

def get_system_metrics(
    db: SessionLocal,
    start_time: Optional[datetime] = None,
    end_time: Optional[datetime] = None,
    limit: int = 100
) -> List[SystemMetric]:
    """Get system metrics with optional time range"""
    query = db.query(SystemMetric)
    if start_time:
        query = query.filter(SystemMetric.timestamp >= start_time)
    if end_time:


    return query.all(), total

def create_task(db: SessionLocal, task_data: Dict[str, Any]) -> Task:
    """Create a new task"""
    task = Task(**task_data)
    db.add(task)
    db.commit()
    db.refresh(task)
    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:
    """Add system metrics"""
    metrics = SystemMetric(**metrics_data)
    db.add(metrics)
    db.commit()
    db.refresh(metrics)
    return metrics

def get_system_metrics(
    db: SessionLocal,
    start_time: Optional[datetime] = None,
    end_time: Optional[datetime] = None,
    limit: int = 100
) -> List[SystemMetric]:
    """Get system metrics with optional time range"""
    query = db.query(SystemMetric)
    if start_time:
        query = query.filter(SystemMetric.timestamp >= start_time)
    if end_time:
        query = query.filter(SystemMetric.timestamp <= end_time)
    return query.order_by(SystemMetric.timestamp.desc()).limit(limit).all()


    return task

def update_task(db: SessionLocal, task_id: str, task_data: Dict[str, Any]) -> Optional[Task]:
    """Update task by ID"""
    task = get_task_by_id(db, task_id)
    if task:
        for key, value in task_data.items():
            setattr(task, key, value)
        db.commit()
        db.refresh(task)
    return task

def add_system_metrics(db: SessionLocal, metrics_data: Dict[str, Any]) -> SystemMetric:
    """Add system metrics"""
    metrics = SystemMetric(**metrics_data)
    db.add(metrics)
    db.commit()
    db.refresh(metrics)
    return metrics

def get_system_metrics(
    db: SessionLocal,
    start_time: Optional[datetime] = None,
    end_time: Optional[datetime] = None,
    limit: int = 100
) -> List[SystemMetric]:
    """Get system metrics with optional time range"""
    query = db.query(SystemMetric)
    if start_time:
        query = query.filter(SystemMetric.timestamp >= start_time)
    if end_time:
        query = query.filter(SystemMetric.timestamp <= end_time)
    return query.order_by(SystemMetric.timestamp.desc()).limit(limit).all()