key_os/server.py

from __future__ import annotations

import logging
import os
import re
import sys
import time
import urllib.parse
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from key_os.facilities import kernel_facilities
from key_os.models import Receipt, ServiceSpec, utc_now
from key_os.probe import ProbeError, call_endpoint, enumerate_space
from key_os.routes import default_links, route_query
from key_os.security import contracts_report, redact, require_confirmation
from key_os.store import KeyOSStore

log = logging.getLogger("key_os")


from pydantic import BaseModel, Field

class ConnectRequest(BaseModel):
    target: str = Field(..., description="Target Space ID (owner/name) or full URL to connect to", json_schema_extra={"example": "tostido/meshscale"})

class ProbeRequest(BaseModel):
    target: str = Field(..., description="Target Space ID or URL", json_schema_extra={"example": "http://127.0.0.1:7861"})
    path: str = Field(..., description="API endpoint path to query", json_schema_extra={"example": "/health"})
    method: str = Field("GET", description="HTTP method to use", json_schema_extra={"example": "GET"})
    query: dict[str, Any] | None = Field(None, description="Optional query parameters")
    body: Any | None = Field(None, description="Optional request body to send as JSON")

class ReceiptRequest(BaseModel):
    event_type: str = Field("operator_note", description="Event category / type", json_schema_extra={"example": "diagnostic_smoke"})
    subject: str = Field("key_os", description="Event subject or target key", json_schema_extra={"example": "local_server"})
    source: str = Field("operator", description="Creator or source identifier", json_schema_extra={"example": "codex_agent"})
    payload: dict[str, Any] = Field(default_factory=dict, description="Custom payload data")

class NotepadAddRequest(BaseModel):
    type: str = Field("observation", description="Note type (observation, hypothesis, action, outcome, question, guide, issue)", json_schema_extra={"example": "issue"})
    title: str | None = Field(None, description="Title of the note (defaults to type if omitted)")
    content: str = Field(..., description="The main text/content of the note", json_schema_extra={"example": "Local probes fail due to HF host checks."})
    tags: list[str] = Field(default_factory=list, description="Tags to organize the notes")
    links: list[str] = Field(default_factory=list, description="Associated links or file paths")
    confidence: float | None = Field(None, description="Confidence level between 0 and 1")
    source: str = Field("operator", description="Source indicating who added the note", json_schema_extra={"example": "browser"})

class NotepadClearRequest(BaseModel):
    confirm: str = Field(..., description="Set to 'CONFIRM' to authorize deletion", json_schema_extra={"example": "CONFIRM"})

class FactoryResetRequest(BaseModel):
    confirm: str = Field(..., description="Set to 'CONFIRM' to authorize reset", json_schema_extra={"example": "CONFIRM"})
    clear_receipts: bool = Field(True, description="Whether to purge all receipts")
    clear_notepad: bool = Field(False, description="Whether to purge the notepad")
    clear_dispatch: bool = Field(True, description="Whether to clear dispatch state")

class SignalAckRequest(BaseModel):
    state: str = Field(..., description="New acknowledgement state (new, seen, deferred, handled, blocked, superseded)", json_schema_extra={"example": "handled"})
    note: str | None = Field(None, description="Optional comment regarding the status change", json_schema_extra={"example": "fixed in 0.1.3"})

class ContinuityCheckpointRequest(BaseModel):
    subject: str = Field("codex_continuity", description="Optional subject annotation")
    note: str = Field("Codex continuity checked.", description="Optional note text")
    source: str = Field("operator", description="Optional source label")

class DispatchStartRequest(BaseModel):
    seed: str = Field("status", description="The initial dispatch seed")
    allow_mutating: bool = Field(True, description="Whether to permit mutating routes")
    allow_external: bool = Field(True, description="Whether to permit external HTTP routes")
    allow_destructive: bool = Field(True, description="Whether to permit destructive routes")
    max_steps: int = Field(60, description="Max execution steps allowed")
    step_delay: float = Field(0.45, description="Delay between execution steps in seconds")


ROOT = Path(__file__).parent
STATIC = ROOT / "static"
SERVICE_KEY_RE = re.compile(r"^[a-z0-9][a-z0-9_.-]{0,80}$")
RESERVED_SERVICE_KEYS = frozenset(facility.key for facility in kernel_facilities())


def _bool_env(name: str, default: bool = False) -> bool:
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _store() -> KeyOSStore:
    return KeyOSStore(os.environ.get("KEY_OS_ROOT", "/data/key_os"))


def _links() -> dict[str, str]:
    links = default_links()
    if os.environ.get("KEY_OS_HOME_URL"):
        links["home"] = os.environ["KEY_OS_HOME_URL"]
    if os.environ.get("KEY_OS_BUCKET_URL"):
        links["bucket"] = os.environ["KEY_OS_BUCKET_URL"]
    if os.environ.get("KEY_OS_DOCS_URL"):
        links["docs"] = os.environ["KEY_OS_DOCS_URL"]
    return links


def _runtime_flags() -> dict[str, Any]:
    return {
        "operator_auth_required": False,
        "operator_auth_configured": False,
        "unauth_mutation_allowed": True,
    }


def _boot_diagnostics() -> None:
    """Log a clear summary of storage config at startup."""
    root = os.environ.get("KEY_OS_ROOT", "/data/key_os")
    log.info("──────────────────────────────────────────────")
    log.info("  KEY OS boot diagnostics")
    log.info("──────────────────────────────────────────────")
    log.info(
        "  ROOT            : %s (writable=%s)",
        root,
        os.access(root, os.W_OK) if os.path.isdir(root) else "dir-missing",
    )
    log.info("  AUTH            : Wholesale Removed (Open Mutability)")
    log.info("──────────────────────────────────────────────")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup/shutdown lifecycle — runs after env is fully available."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(levelname)s:     %(name)s - %(message)s",
        stream=sys.stderr,
    )
    _boot_diagnostics()
    yield


app = FastAPI(
    title=os.environ.get("KEY_OS_NAME", "KEY OS"),
    lifespan=lifespan,
)
app.mount("/static", StaticFiles(directory=str(STATIC)), name="static")


@app.get("/")
async def index() -> FileResponse:
    return FileResponse(STATIC / "index.html")


@app.get("/health")
async def health() -> dict[str, Any]:
    return {"ok": True, "surface": "key_os", "time": utc_now()}


@app.get("/api/status")
async def status(detail: str = "compact") -> dict[str, Any]:
    return _store().status(_links(), runtime_flags=_runtime_flags(), detail=detail)


@app.get("/api/os")
async def resource_os() -> dict[str, Any]:
    return _store().resource_snapshot(_links(), runtime_flags=_runtime_flags())


@app.get("/api/os/search")
async def resource_search(q: str = "", limit: int = 40) -> dict[str, Any]:
    return _store().resource_search(q, limit=limit, links=_links(), runtime_flags=_runtime_flags())


@app.get("/api/os/oculus")
async def oculus() -> dict[str, Any]:
    return _store().status(_links(), runtime_flags=_runtime_flags()).get("oculus", {})


@app.get("/api/continuity")
async def continuity() -> dict[str, Any]:
    return _store().continuity()


@app.get("/api/signals")
async def signals() -> dict[str, Any]:
    return _store().signals()


@app.get("/api/signals/{signal_id}")
async def signal_detail(signal_id: str) -> dict[str, Any]:
    signal = _store().get_signal(signal_id)
    if signal is None:
        raise HTTPException(status_code=404, detail="signal not found")
    return signal


@app.post("/api/signals/{signal_id}/ack")
async def signal_ack(signal_id: str, payload: SignalAckRequest) -> dict[str, Any]:
    try:
        return _store().acknowledge_signal(signal_id, payload.model_dump())
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="signal not found") from exc
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.post("/api/continuity/checkpoint")
async def continuity_checkpoint(payload: ContinuityCheckpointRequest) -> dict[str, Any]:
    return _store().acknowledge_continuity(payload.model_dump())


@app.get("/api/services")
async def services() -> dict[str, Any]:
    custom_count = len(_store().load_services(include_kernel=False))
    items = [service.to_dict(redact_payload=True) for service in _store().load_services()]
    groups: dict[str, list[dict[str, Any]]] = {}
    for item in items:
        groups.setdefault(str(item.get("group") or "custom"), []).append(item)
    return {"count": len(items), "custom_count": custom_count, "groups": groups, "services": items}


@app.get("/api/facilities")
async def facilities() -> dict[str, Any]:
    return _store().registry().manifest()


@app.get("/api/facilities/graph")
async def facility_graph() -> dict[str, Any]:
    return _store().registry().graph()


@app.get("/api/facility/{key}")
async def facility_detail(key: str) -> dict[str, Any]:
    return _store().registry().node_detail(key)


@app.post("/api/services")
async def register_service(request: Request) -> dict[str, Any]:
    payload = await request.json()
    key = str(payload.get("key") or "").strip().lower()
    label = str(payload.get("label") or key).strip()
    group = str(payload.get("group") or "custom").strip().lower()
    route = str(payload.get("route") or "").strip()
    if not key:
        raise HTTPException(status_code=400, detail="service key is required")
    if not SERVICE_KEY_RE.match(key):
        raise HTTPException(status_code=400, detail="service key must match ^[a-z0-9][a-z0-9_.-]{0,80}$")
    if key in RESERVED_SERVICE_KEYS:
        raise HTTPException(status_code=409, detail=f"service key '{key}' is reserved by the kernel")
    if group == "kernel":
        raise HTTPException(status_code=400, detail="custom services may not use reserved group 'kernel'")
    if not route:
        raise HTTPException(status_code=400, detail="service route is required")
    service = ServiceSpec(
        key=key,
        label=label,
        group=group or "custom",
        description=str(payload.get("description") or ""),
        route=route,
        method=str(payload.get("method") or "GET").upper(),
        risk_class=str(payload.get("risk_class") or "inspect"),
        tags=tuple(payload.get("tags") or ()),
        input_schema=payload.get("input_schema") or {},
        output_schema=payload.get("output_schema") or {},
        payload=payload.get("payload") or {},
    )
    return {"ok": True, "service": _store().put_service(service).to_dict(redact_payload=True)}


@app.post("/api/route")
async def route(request: Request) -> dict[str, Any]:
    payload = await request.json()
    query = str(payload.get("query") or "")
    report = _store().status(_links(), runtime_flags=_runtime_flags())
    return {"query": query, "route": route_query(query, report, _links()), "summary": report.get("kernel", {})}


def _is_self_target(url: str, request: Request) -> bool:
    try:
        parsed_target = urllib.parse.urlparse(url)
        target_host = parsed_target.hostname or ""
        target_port = parsed_target.port
        
        parsed_self = urllib.parse.urlparse(str(request.base_url))
        self_host = parsed_self.hostname or ""
        self_port = parsed_self.port or 80
        
        # Normalize localhost/127.0.0.1/::1/127.*
        is_target_local = target_host.lower() in {"localhost", "127.0.0.1", "::1"} or target_host.startswith("127.")
        is_self_local = self_host.lower() in {"localhost", "127.0.0.1", "::1"} or self_host.startswith("127.")
        
        if is_target_local and is_self_local:
            return (target_port or 80) == self_port
            
        return target_host.lower() == self_host.lower() and (target_port or 80) == self_port
    except Exception:
        return False


@app.post("/api/connect")
async def connect_space(payload: ConnectRequest, request: Request) -> dict[str, Any]:
    """Surgeon's tools: quantify a parallel Space's entire API array. No token —
    this is an observation surface, scoped to the HF constellation in probe.py."""
    target = payload.target.strip()
    if not target:
        raise HTTPException(status_code=400, detail="target space url or id is required")
    
    try:
        from key_os.probe import normalize_target
        normalized_target = normalize_target(target)
    except ProbeError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    if _is_self_target(normalized_target, request):
        import httpx
        try:
            transport = httpx.ASGITransport(app=app)
            async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
                t0 = time.perf_counter()
                health_resp = await client.get("/health")
                elapsed = round((time.perf_counter() - t0) * 1000, 1)
                
                spec_resp = await client.get("/openapi.json")
                spec = spec_resp.json() if spec_resp.status_code == 200 else None
                
            if spec and "paths" in spec:
                from key_os.probe import _parse_openapi
                endpoints = _parse_openapi(spec)
                info = spec.get("info") or {}
                result = {
                    "target": target,
                    "base_url": normalized_target,
                    "fetched_at": utc_now(),
                    "health": {
                        "reachable": True,
                        "probe_path": "/health",
                        "status": health_resp.status_code,
                        "stage": "ok" if health_resp.status_code == 200 else "error",
                        "elapsed_ms": elapsed,
                    },
                    "status": "connected",
                    "title": info.get("title"),
                    "version": info.get("version"),
                    "openapi": spec.get("openapi"),
                    "count": len(endpoints),
                    "endpoints": endpoints,
                }
            else:
                result = {
                    "target": target,
                    "base_url": normalized_target,
                    "fetched_at": utc_now(),
                    "health": {
                        "reachable": True,
                        "probe_path": "/health",
                        "status": health_resp.status_code,
                        "stage": "ok" if health_resp.status_code == 200 else "error",
                        "elapsed_ms": elapsed,
                    },
                    "status": "no_openapi",
                    "raw_status": spec_resp.status_code,
                    "message": f"No OpenAPI spec at {normalized_target}/openapi.json",
                    "endpoints": [],
                }
        except Exception as exc:
            result = {
                "target": target,
                "base_url": normalized_target,
                "fetched_at": utc_now(),
                "health": {"reachable": False, "stage": "unreachable", "status": 0},
                "status": "unreachable",
                "message": f"Self-connection error: {exc}",
                "endpoints": [],
            }
    else:
        try:
            result = enumerate_space(target)
        except ProbeError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    try:
        _store().append_event_receipt(
            "space_connect",
            str(result.get("base_url") or target),
            {"status": result.get("status"), "count": result.get("count")},
            source="probe",
        )
    except Exception:  # provenance is best-effort; never block the observation
        pass
    return result


@app.post("/api/probe")
async def probe_space(payload: ProbeRequest, request: Request) -> dict[str, Any]:
    """Fire a single request at a connected Space and return the raw result."""
    target = payload.target.strip()
    path = payload.path.strip()
    method = payload.method.upper()
    if not target or not path:
        raise HTTPException(status_code=400, detail="target and path are required")
    
    try:
        from key_os.probe import normalize_target
        normalized_target = normalize_target(target)
    except ProbeError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    if _is_self_target(normalized_target, request):
        import httpx
        if not path.startswith("/"):
            path = "/" + path
        try:
            transport = httpx.ASGITransport(app=app)
            async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
                t0 = time.perf_counter()
                resp = await client.request(
                    method=method,
                    url=path,
                    params=payload.query,
                    json=payload.body,
                    timeout=10.0,
                )
                elapsed = round((time.perf_counter() - t0) * 1000, 1)
                
                try:
                    body_json = resp.json()
                except Exception:
                    body_json = None
                
                result = {
                    "ok": 200 <= resp.status_code < 400,
                    "status": resp.status_code,
                    "method": method,
                    "url": normalized_target + path,
                    "elapsed_ms": elapsed,
                    "content_type": resp.headers.get("content-type", ""),
                    "json": body_json,
                    "text": None if body_json is not None else resp.text[:20000],
                    "truncated": len(resp.text) > 20000,
                }
        except Exception as exc:
            result = {
                "ok": False,
                "status": 0,
                "method": method,
                "url": normalized_target + path,
                "error": f"{type(exc).__name__}: {exc}",
                "elapsed_ms": 0.0,
            }
    else:
        try:
            result = call_endpoint(
                target,
                path,
                method=method,
                query=payload.query,
                body=payload.body,
            )
        except ProbeError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    try:
        _store().append_event_receipt(
            "space_probe",
            f"{method} {path}",
            {"target": target, "status": result.get("status"), "elapsed_ms": result.get("elapsed_ms")},
            source="probe",
        )
    except Exception:
        pass
    return result


@app.get("/api/receipts")
async def receipts(limit: int = 50) -> dict[str, Any]:
    items = _store().load_receipts(limit=limit)
    return {"count": len(items), "receipts": items}


@app.post("/api/receipts")
async def write_receipt(payload: ReceiptRequest) -> dict[str, Any]:
    receipt = Receipt(
        event_type=payload.event_type,
        subject=payload.subject,
        source=payload.source,
        payload=redact(payload.payload),
    )
    _store().append_receipt(receipt)
    return {"ok": True, "receipt": receipt.to_dict()}


@app.get("/api/notepad")
async def notepad(query: str = "", limit: int = 80, type: str = "") -> dict[str, Any]:
    items = _store().notepad().list(query=query, limit=limit, note_type=type)
    return {"schema": "key_os.research_notepad/v1", "count": len(items), "entries": items}


@app.get("/api/notepad/summary")
async def notepad_summary() -> dict[str, Any]:
    return _store().notepad().summary()


@app.post("/api/notepad")
async def notepad_add(payload: NotepadAddRequest) -> dict[str, Any]:
    entry = _store().notepad().add(payload.model_dump(), source=payload.source)
    receipt = _store().append_event_receipt("notepad_add", str(entry.get("id")), entry, source="notepad")
    return {"ok": True, "entry": entry, "receipt": receipt.to_dict()}


@app.post("/api/notepad/clear")
async def notepad_clear(payload: NotepadClearRequest) -> dict[str, Any]:
    try:
        require_confirmation("notepad_clear", payload.model_dump())
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    result = _store().notepad().clear()
    receipt = _store().append_event_receipt("notepad_clear", "research_notepad", result, source="notepad")
    return {**result, "receipt": receipt.to_dict()}


@app.get("/api/dispatch/state")
async def dispatch_state() -> dict[str, Any]:
    store = _store()
    state = store.dispatch().read()
    state["plan"] = store.dispatch().plan(store.registry(), policy=state.get("policy") or {})
    return state


@app.post("/api/dispatch/start")
async def dispatch_start(payload: DispatchStartRequest) -> dict[str, Any]:
    store = _store()
    state = store.dispatch().start(payload.model_dump(), store.registry())
    receipt = store.append_event_receipt("dispatch_start", str(payload.seed), payload.model_dump(), source="dispatch")
    return {"ok": True, "state": state, "receipt": receipt.to_dict()}


@app.post("/api/dispatch/stop")
async def dispatch_stop(request: Request) -> dict[str, Any]:
    store = _store()
    state = store.dispatch().stop()
    receipt = store.append_event_receipt("dispatch_stop", "dispatcher", {}, source="dispatch")
    return {"ok": True, "state": state, "receipt": receipt.to_dict()}


@app.get("/api/security/contracts")
async def security_contracts() -> dict[str, Any]:
    return contracts_report()


@app.post("/api/factory-reset")
async def factory_reset(payload: FactoryResetRequest) -> dict[str, Any]:
    try:
        require_confirmation("factory_reset", payload.model_dump())
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    clear_receipts = payload.clear_receipts
    clear_notepad = payload.clear_notepad
    clear_dispatch = payload.clear_dispatch
    return _store().factory_reset(clear_receipts=clear_receipts, clear_notepad=clear_notepad, clear_dispatch=clear_dispatch)


@app.post("/api/reset")
async def legacy_factory_reset(payload: FactoryResetRequest) -> dict[str, Any]:
    return await factory_reset(payload)