app.py

import io
import os
import json
import time
import uuid
import asyncio
import urllib.request
import urllib.error

from fastapi import FastAPI, HTTPException, Request, Response, Depends
from fastapi.responses import StreamingResponse, FileResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from PIL import Image

from seed_catalog import SeedCatalog
from breeder import breed, clone
from starter_loader import StarterSeedLoader
from name_generator import NameGenerator
from icon_loader import generate_seed_image
from seed_profile import GrowthStage
import lab
import deck

# Embed the K-os Observer Bus in-process: WEED-SIM and the mounted bus HUD share
# one file-backed store, so breeding receipts and the signal inbox are first-class.
os.environ.setdefault("KEY_OS_ROOT", os.path.join(os.path.dirname(__file__), ".observer-bus-runtime"))
os.environ.setdefault("KEY_OS_NAME", "WEED-SIM Observer Bus")
import keyos_bridge

app = FastAPI(title="WEED-SIM Web API")

# Mount the full K-os Observer Bus (API + HUD) under /bus, sharing KEY_OS_ROOT.
try:
    from key_os.server import app as kos_app
    app.mount("/bus", kos_app)
except Exception as _kos_exc:  # never let the bus block the game from starting
    print(f"[K-os] bus mount skipped: {_kos_exc}")


def _register_facilities():
    """Advertise WEED-SIM's operations as K-os facilities so the bus exposes a
    callable syscall table that agents can discover at /bus/api/facilities."""
    try:
        from key_os.store import KeyOSStore
        from key_os.models import ServiceSpec
        store = KeyOSStore(os.environ["KEY_OS_ROOT"])
        existing = {s.key for s in store.load_services(include_kernel=False)}
        specs = [
            ServiceSpec(key="weedsim.inventory", label="Inventory", group="weed-sim",
                        description="List the session's specimens with traits + colors.",
                        route="/api/inventory", method="GET", risk_class="inspect", tags=("weed-sim", "read")),
            ServiceSpec(key="weedsim.load_starters", label="Load Starters", group="weed-sim",
                        description="Seed the session with the 12 baseline strains.",
                        route="/api/inventory/starter", method="POST", risk_class="mutating", tags=("weed-sim",)),
            ServiceSpec(key="weedsim.grow", label="Grow", group="weed-sim",
                        description="Advance a specimen's growth stage.",
                        route="/api/seed/{id}/grow", method="POST", risk_class="mutating", tags=("weed-sim",)),
            ServiceSpec(key="weedsim.breed", label="Breed", group="weed-sim",
                        description="Cross two parents; HSV recombination. Emits a provenance receipt.",
                        route="/api/breed", method="POST", risk_class="mutating", tags=("weed-sim", "genetics")),
            ServiceSpec(key="weedsim.clone", label="Clone", group="weed-sim",
                        description="Vegetative copy of one specimen.",
                        route="/api/clone", method="POST", risk_class="mutating", tags=("weed-sim", "genetics")),
            ServiceSpec(key="weedsim.germplasm", label="Germplasm", group="weed-sim",
                        description="BrAPI-style accessions: generation, pedigree, observations, genotype, facets.",
                        route="/api/lab/germplasm", method="GET", risk_class="inspect", tags=("weed-sim", "brapi")),
            ServiceSpec(key="weedsim.predict_cross", label="Predict Cross", group="weed-sim",
                        description="Monte-Carlo offspring trait + color distributions for two parents.",
                        route="/api/lab/predict-cross", method="POST", risk_class="inspect", tags=("weed-sim", "analysis")),
            ServiceSpec(key="weedsim.newick", label="Newick Export", group="weed-sim",
                        description="Reticulate pedigree as Extended-Newick.",
                        route="/api/lab/newick", method="GET", risk_class="inspect", tags=("weed-sim", "phylo")),
        ]
        for spec in specs:
            if spec.key not in existing:
                store.put_service(spec)
    except Exception as exc:
        print(f"[K-os] facility registration skipped: {exc}")


_register_facilities()

# NameGenerator is stateless, so a single shared instance is fine.
name_gen = NameGenerator()

# ── Per-visitor session state ────────────────────────────────────────────────
# A public Space is multi-user. Keeping one global catalog means every visitor
# shares (and can wipe) the same inventory. Instead each browser gets its own
# cookie-keyed catalog, held in memory and bounded so the dict can't grow
# without limit. State is intentionally ephemeral; it resets when the Space
# restarts, which is the right behaviour for a sandbox game demo.
SESSION_COOKIE = "weedsim_sid"
SESSION_HEADER = "x-weedsim-session"   # primary: survives cross-site iframes (no cookie needed)
SESSION_TTL = 60 * 60 * 6          # 6 hours of inactivity
MAX_SESSIONS = 500                 # hard cap on concurrent sessions
_sessions: dict[str, dict] = {}


def _new_session() -> dict:
    catalog = SeedCatalog()
    return {
        "catalog": catalog,
        "loader": StarterSeedLoader(catalog, name_gen),
        "selection": [],   # shared specimen selection across Play / Lab / Deck
        "ts": time.time(),
    }


def _reap_sessions() -> None:
    """Drop expired sessions, then evict oldest if still over the cap."""
    now = time.time()
    for sid in [s for s, v in _sessions.items() if now - v["ts"] > SESSION_TTL]:
        _sessions.pop(sid, None)
    if len(_sessions) > MAX_SESSIONS:
        for sid, _ in sorted(_sessions.items(), key=lambda kv: kv[1]["ts"])[
            : len(_sessions) - MAX_SESSIONS
        ]:
            _sessions.pop(sid, None)


def get_session(request: Request, response: Response) -> dict:
    """Resolve (or create) the caller's session.

    Identity comes from (in order) the X-WeedSim-Session header (set by the
    frontend from a localStorage token — works in any iframe), a `sid` query
    param (so <img> sprite URLs carry it), or the cookie (fallback). This makes
    sessions independent of third-party-cookie policy.
    """
    sid = (
        request.headers.get(SESSION_HEADER)
        or request.query_params.get("sid")
        or request.cookies.get(SESSION_COOKIE)
    )
    if not sid or sid not in _sessions:
        _reap_sessions()
        if not sid:
            sid = uuid.uuid4().hex
        _sessions[sid] = _new_session()
        response.set_cookie(
            SESSION_COOKIE, sid, httponly=True, samesite="none", secure=True, max_age=SESSION_TTL
        )
    sess = _sessions[sid]
    sess["ts"] = time.time()
    return sess


def lookup_session(request: Request) -> dict | None:
    """Read-only session resolution for image endpoints (never sets a cookie)."""
    sid = request.cookies.get(SESSION_COOKIE)
    sess = _sessions.get(sid) if sid else None
    if sess:
        sess["ts"] = time.time()
    return sess


# ── Static frontend ──────────────────────────────────────────────────────────
app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/")
async def root():
    return FileResponse("static/landing.html")


@app.get("/play")
async def play_page():
    return FileResponse("static/index.html")


@app.get("/docs-hub")
async def docs_page():
    return FileResponse("static/docs.html")


@app.get("/healthz")
async def healthz():
    return {"status": "ok", "sessions": len(_sessions)}


# Meta-routing manifest: every observable surface, so the doc hub can link
# straight into the live system (a facility table, K-os style).
ROUTE_MANIFEST = {
    "schema": "weedsim.route_manifest/v1",
    "groups": [
        {"group": "Surfaces", "routes": [
            {"method": "GET", "path": "/", "label": "Landing", "kind": "page", "desc": "Meta entry + live Observer-Bus state."},
            {"method": "GET", "path": "/play", "label": "Play Mode", "kind": "page", "desc": "Grow, breed, and clone specimens."},
            {"method": "GET", "path": "/lab", "label": "Speciation Lab", "kind": "page", "desc": "Faceted grid, parallel coordinates, gamut wheel, cross predictor."},
            {"method": "GET", "path": "/deck", "label": "The Deck", "kind": "page", "desc": "Shared-state hub: gauges, breakdowns, recent crosses, signals."},
            {"method": "GET", "path": "/agent", "label": "Agent", "kind": "page", "desc": "Chat agent that drives the whole system via HF Inference Providers + tool calls."},
            {"method": "GET", "path": "/docs", "label": "OpenAPI (Swagger)", "kind": "page", "desc": "Auto-generated interactive API explorer."},
            {"method": "GET", "path": "/bus", "label": "Observer Bus (embedded)", "kind": "page", "desc": "The full K-os HUD running in-process: signals, receipts, facilities, Resource OS."},
        ]},
        {"group": "Play API", "routes": [
            {"method": "GET", "path": "/api/inventory", "label": "Inventory", "kind": "get", "desc": "Session specimens with bud/leaf colors + traits."},
            {"method": "POST", "path": "/api/inventory/starter", "label": "Load starters", "kind": "post", "desc": "Add the 12 baseline strains."},
            {"method": "POST", "path": "/api/inventory/clear", "label": "Clear", "kind": "post", "desc": "Empty the session registry."},
            {"method": "POST", "path": "/api/seed/{id}/grow", "label": "Grow", "kind": "post", "desc": "Advance a specimen's growth stage."},
            {"method": "POST", "path": "/api/breed", "label": "Breed", "kind": "post", "desc": "Cross two parents (HSV recombination); emits a K-os receipt."},
            {"method": "POST", "path": "/api/clone", "label": "Clone", "kind": "post", "desc": "Vegetative copy of one specimen."},
            {"method": "GET", "path": "/api/image/{id}/{size}", "label": "Sprite", "kind": "get-tmpl", "desc": "Pixel-art sprite, recolored via bud/leaf masks."},
        ]},
        {"group": "Lab API (BrAPI-style)", "routes": [
            {"method": "GET", "path": "/api/lab/germplasm", "label": "Germplasm", "kind": "get", "desc": "Accessions: generation, pedigree, observations, genotype, hue, facets."},
            {"method": "POST", "path": "/api/lab/predict-cross", "label": "Predict cross", "kind": "post", "desc": "Monte-Carlo offspring trait + color distributions."},
            {"method": "GET", "path": "/api/lab/newick", "label": "Newick export", "kind": "get", "desc": "Reticulate pedigree as Extended-Newick."},
        ]},
        {"group": "Deck API", "routes": [
            {"method": "GET", "path": "/api/deck/summary", "label": "Summary", "kind": "get", "desc": "Unified state: counts, generations, trait averages, gamut, recent crosses."},
            {"method": "GET", "path": "/api/deck/signals", "label": "Signals", "kind": "get", "desc": "Observer-Bus signal inbox (fail-open proxy to K-os)."},
            {"method": "GET", "path": "/api/selection", "label": "Selection", "kind": "get", "desc": "Shared specimen selection across modes."},
            {"method": "GET", "path": "/api/routes", "label": "Route manifest", "kind": "get", "desc": "This manifest."},
            {"method": "GET", "path": "/healthz", "label": "Health", "kind": "get", "desc": "Liveness + active session count."},
        ]},
        {"group": "Observer Bus API (embedded, agent-usable)", "routes": [
            {"method": "GET", "path": "/bus/api/signals", "label": "Bus signals", "kind": "get", "desc": "The live signal inbox (crosses appear as provenance receipts)."},
            {"method": "GET", "path": "/bus/api/os", "label": "Resource OS", "kind": "get", "desc": "Read-only packet map over facilities, receipts, notepad, boundaries."},
            {"method": "GET", "path": "/bus/api/os/search?q=", "label": "Resource search", "kind": "get", "desc": "Search Resource-OS packets."},
            {"method": "GET", "path": "/bus/api/facilities", "label": "Facilities", "kind": "get", "desc": "Callable syscall table — WEED-SIM ops registered as facilities."},
            {"method": "POST", "path": "/bus/api/receipts", "label": "Publish receipt", "kind": "post", "desc": "Append a provenance receipt (no auth on this build)."},
        ]},
        {"group": "External / Provenance", "routes": [
            {"method": "GET", "path": "https://huggingface.co/spaces/tostido/K-os", "label": "K-os Observer Bus", "kind": "external", "desc": "The OS substrate WEED-SIM publishes receipts to."},
            {"method": "GET", "path": "https://brapi.org/specification", "label": "BrAPI spec", "kind": "external", "desc": "Breeding API the germplasm model follows."},
            {"method": "GET", "path": "https://icytree.org/", "label": "IcyTree", "kind": "external", "desc": "Open a Newick export to view the pedigree network."},
        ]},
    ],
}


@app.get("/api/routes")
async def route_manifest():
    return ROUTE_MANIFEST


# ── Agent chat: proxy to OpenAI-compatible inference routers ──────────────────
# Two providers: Hugging Face Inference Providers and OpenRouter. The visitor's
# token is passed per-request via X-HF-Token and NEVER stored — this proxy just
# forwards (dodging browser CORS + keeping the token off our server). The agent
# loop itself runs in the browser.
PROVIDER_BASE = {
    "hf": "https://router.huggingface.co/v1",
    "openrouter": "https://openrouter.ai/api/v1",
}


def _provider_token(request: Request, provider: str) -> str:
    tok = request.headers.get("x-hf-token")  # the UI's single "provider token" field
    if tok:
        return tok
    return os.environ.get("OPENROUTER_API_KEY", "") if provider == "openrouter" else os.environ.get("HF_TOKEN", "")


def _extra_headers(provider: str) -> dict:
    if provider == "openrouter":
        return {"HTTP-Referer": "https://tostido-weed-sim.hf.space", "X-Title": "WEED-SIM Agent"}
    return {}


def _router_get(base: str, path: str, token: str, extra: dict | None = None):
    headers = {"Authorization": f"Bearer {token}", **(extra or {})}
    req = urllib.request.Request(base + path, method="GET", headers=headers)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read())


def _router_post(base: str, path: str, token: str, payload: dict, extra: dict | None = None):
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json", **(extra or {})}
    req = urllib.request.Request(base + path, data=json.dumps(payload).encode(), method="POST", headers=headers)
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.loads(r.read())


def _normalize_models(provider: str, data) -> list:
    items = data.get("data", data if isinstance(data, list) else [])
    out = []
    for m in items:
        if not isinstance(m, dict):
            continue
        if provider == "openrouter":
            sp = m.get("supported_parameters") or []
            pricing = m.get("pricing") or {}
            free = str(pricing.get("prompt", "0")) in ("0", "0.0") or ":free" in str(m.get("id", ""))
            out.append({"id": m.get("id"), "tools": "tools" in sp, "free": free,
                        "context": int(m.get("context_length") or 0), "providers": []})
        else:
            provs = [p for p in (m.get("providers") or []) if isinstance(p, dict)]
            ctx = max([int(p.get("context_length") or 0) for p in provs] or [0])
            out.append({"id": m.get("id"), "tools": any(p.get("supports_tools") for p in provs),
                        "free": any(p.get("is_free") for p in provs), "context": ctx,
                        "providers": [p.get("provider") for p in provs]})
    out.sort(key=lambda x: (not x["tools"], -(x["context"] or 0), str(x["id"]).lower()))
    return out


@app.get("/agent")
async def agent_page():
    return FileResponse("static/agent.html")


@app.get("/api/agent/models")
async def agent_models(request: Request, provider: str = "hf"):
    base = PROVIDER_BASE.get(provider, PROVIDER_BASE["hf"])
    token = _provider_token(request, provider)
    if not token:
        raise HTTPException(status_code=401, detail="Provider token required.")
    try:
        data = await asyncio.to_thread(_router_get, base, "/models", token, _extra_headers(provider))
    except urllib.error.HTTPError as e:
        raise HTTPException(status_code=e.code, detail=e.read().decode()[:300])
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))
    models = _normalize_models(provider, data)
    return {"count": len(models), "provider": provider, "models": models}


class AgentChatRequest(BaseModel):
    model: str
    messages: list
    provider: str = "hf"
    tools: list | None = None
    tool_choice: str | None = "auto"
    temperature: float | None = 0.4
    max_tokens: int | None = 1200


@app.post("/api/agent/chat")
async def agent_chat(req: AgentChatRequest, request: Request):
    base = PROVIDER_BASE.get(req.provider, PROVIDER_BASE["hf"])
    token = _provider_token(request, req.provider)
    if not token:
        raise HTTPException(status_code=401, detail="Paste a provider token to use the agent.")
    payload = {"model": req.model, "messages": req.messages}
    if req.tools:
        payload["tools"] = req.tools
        payload["tool_choice"] = req.tool_choice or "auto"
    if req.temperature is not None:
        payload["temperature"] = req.temperature
    if req.max_tokens:
        payload["max_tokens"] = req.max_tokens
    try:
        return await asyncio.to_thread(_router_post, base, "/chat/completions", token, payload, _extra_headers(req.provider))
    except urllib.error.HTTPError as e:
        raise HTTPException(status_code=e.code, detail=f"{req.provider} router: {e.read().decode()[:400]}")
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"{req.provider} router unreachable: {e}")


@app.get("/lab")
async def lab_page():
    return FileResponse("static/lab.html")


@app.get("/deck")
async def deck_page():
    return FileResponse("static/deck.html")


# ── The Deck: shared-state hub ────────────────────────────────────────────────
@app.get("/api/deck/summary")
async def deck_summary(session: dict = Depends(get_session)):
    return deck.summary(session["catalog"].get_all_seeds())


@app.get("/api/deck/signals")
async def deck_signals():
    return keyos_bridge.signals()


class SelectionRequest(BaseModel):
    ids: list[str] = []


@app.get("/api/selection")
async def get_selection(session: dict = Depends(get_session)):
    return {"ids": session.get("selection", [])}


@app.post("/api/selection")
async def set_selection(req: SelectionRequest, session: dict = Depends(get_session)):
    catalog = session["catalog"]
    session["selection"] = [i for i in req.ids if catalog.get_seed_by_id(i)][:8]
    return {"ids": session["selection"]}


# ── Science Lab (BrAPI-style germplasm + speciation surface) ──────────────────
@app.get("/api/lab/germplasm")
async def lab_germplasm(session: dict = Depends(get_session)):
    return lab.lab_payload(session["catalog"].get_all_seeds())


class CrossRequest(BaseModel):
    parent1_id: str
    parent2_id: str
    n: int = 200


@app.post("/api/lab/predict-cross")
async def lab_predict_cross(req: CrossRequest, session: dict = Depends(get_session)):
    catalog = session["catalog"]
    p1 = catalog.get_seed_by_id(req.parent1_id)
    p2 = catalog.get_seed_by_id(req.parent2_id)
    if not p1 or not p2:
        raise HTTPException(status_code=404, detail="Parents not found")
    return lab.predict_cross(p1, p2, n=max(20, min(1000, req.n)))


@app.get("/api/lab/newick")
async def lab_newick(session: dict = Depends(get_session)):
    return PlainTextResponse(lab.to_newick(session["catalog"].get_all_seeds()))


# ── Inventory ────────────────────────────────────────────────────────────────
@app.get("/api/inventory")
async def get_inventory(session: dict = Depends(get_session)):
    out = []
    for s in session["catalog"].get_all_seeds():
        out.append({
            "id": s.seed_id,
            "name": s.strain_name,
            "stage": s.growth_stage.name,
            "type": s.type,
            "thc": s.thc,
            "cbd": s.cbd,
            "yield": s.yield_,
            "grow_time": s.grow_time,
            "bud_color": list(s.bud_color),
            "leaf_color": list(s.leaf_color),
            "stabilities": s.stabilities,
            "alleles": s.alleles,
            "lineage": list(s.lineage),
            "can_attempt": s.can_attempt(),
            "attempts_used": s.attempts_used,
            "max_attempts": s.max_attempts,
            "is_starter": s.is_starter,
            "description": s.description,
            "quantity": s.quantity,
        })
    return {"seeds": out}


@app.post("/api/inventory/starter")
async def load_starters(session: dict = Depends(get_session)):
    session["loader"].load_starter_seeds()
    return {"ok": True}


@app.post("/api/inventory/clear")
async def clear_inventory(session: dict = Depends(get_session)):
    session["catalog"].clear()
    return {"ok": True}


@app.post("/api/seed/{seed_id}/grow")
async def grow_seed(seed_id: str, session: dict = Depends(get_session)):
    seed = session["catalog"].get_seed_by_id(seed_id)
    if not seed:
        raise HTTPException(status_code=404, detail="Seed not found")
    seed.advance_growth_stage()
    return {"ok": True}


# ── Breeding / cloning ───────────────────────────────────────────────────────
class BreedRequest(BaseModel):
    parent1_id: str
    parent2_id: str


@app.post("/api/breed")
async def breed_seeds(req: BreedRequest, session: dict = Depends(get_session)):
    catalog = session["catalog"]
    p1 = catalog.get_seed_by_id(req.parent1_id)
    p2 = catalog.get_seed_by_id(req.parent2_id)
    if not p1 or not p2:
        raise HTTPException(status_code=404, detail="Parents not found")
    try:
        offspring = breed(p1, p2, name_gen)
        catalog.add_seed(offspring)
        keyos_bridge.publish_receipt("cross", offspring.strain_name, {
            "offspring_id": offspring.seed_id,
            "parents": [p1.strain_name, p2.strain_name],
            "bud_color": list(offspring.bud_color),
            "leaf_color": list(offspring.leaf_color),
        })
        return {"ok": True, "offspring_id": offspring.seed_id}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


class CloneRequest(BaseModel):
    seed_id: str


@app.post("/api/clone")
async def clone_seed(req: CloneRequest, session: dict = Depends(get_session)):
    catalog = session["catalog"]
    p = catalog.get_seed_by_id(req.seed_id)
    if not p:
        raise HTTPException(status_code=404, detail="Seed not found")
    try:
        c = clone(p, name_gen)
        catalog.add_seed(c)
        keyos_bridge.publish_receipt("clone", c.strain_name, {
            "clone_id": c.seed_id, "parent": p.strain_name,
        })
        return {"ok": True, "clone_id": c.seed_id}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# ── Sprite rendering ─────────────────────────────────────────────────────────
# Rendered PNGs are deterministic for a given (seed, stage, size), so cache the
# bytes. The grid re-renders on every selection/breed/grow, which would otherwise
# re-run the recolorizer for each visible sprite on every interaction.
_RENDER_CACHE: dict[tuple, bytes] = {}
_RENDER_CACHE_MAX = 2000


def _render(seed, stage_name: str, size: int) -> StreamingResponse:
    key = (seed.seed_id, stage_name, size, str(seed.bud_color), str(seed.leaf_color), seed.is_starter)
    png = _RENDER_CACHE.get(key)
    if png is None:
        img = generate_seed_image(seed, stage_name)
        if size != img.width:
            img = img.resize((size, size), Image.NEAREST)
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        png = buf.getvalue()
        if len(_RENDER_CACHE) >= _RENDER_CACHE_MAX:
            _RENDER_CACHE.pop(next(iter(_RENDER_CACHE)))
        _RENDER_CACHE[key] = png
    return StreamingResponse(io.BytesIO(png), media_type="image/png")


@app.get("/api/image/{seed_id}/{size}")
async def get_seed_image(seed_id: str, size: int, session: dict = Depends(get_session)):
    seed = session["catalog"].get_seed_by_id(seed_id)
    if not seed:
        raise HTTPException(status_code=404, detail="Seed not found")
    return _render(seed, seed.growth_stage.name, size)


@app.get("/api/image_mature/{seed_id}/{size}")
async def get_mature_seed_image(seed_id: str, size: int, session: dict = Depends(get_session)):
    # Forces the mature stage, used for lineage-tree thumbnails.
    seed = session["catalog"].get_seed_by_id(seed_id)
    if not seed:
        raise HTTPException(status_code=404, detail="Seed not found")
    return _render(seed, GrowthStage.MATURE.name, size)


if __name__ == "__main__":
    import uvicorn

    # Hugging Face Spaces (and most PaaS) inject the port via $PORT; default to
    # the HF Docker convention of 7860 for local parity with the container.
    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run("app:app", host="0.0.0.0", port=port)