# Marimo Zero - Official Agent Zero + Marimo + OmniRoute
# Uses frdel/agent-zero-exe base image

FROM frdel/agent-zero-exe:latest

USER root

# 1. Install dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    curl jq python3-pip \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir --break-system-packages marimo fastapi uvicorn requests httpx

# 2. Install OmniRoute
RUN npm install -g omniroute

# 3. Setup directories
WORKDIR /app
RUN mkdir -p /app/marimo /app/workspace /app/.omniroute && chmod -R 777 /app

# 4. Copy Marimo notebook
COPY marimo_sandbox.py /app/marimo/marimo_sandbox.py

# 5. Create startup script
RUN cat > /app/start.sh << 'EOF'
#!/bin/bash
set -e

echo "🚀 Starting Marimo Zero..."

# Parse API keys
if [ -n "$API_KEYS_JSON" ]; then
    echo "📝 Loading API keys..."
    export OPENAI_API_KEY=$(echo $API_KEYS_JSON | jq -r '.openai // empty')
    export ANTHROPIC_API_KEY=$(echo $API_KEYS_JSON | jq -r '.anthropic // empty')
fi

# Start OmniRoute
echo "📡 Starting OmniRoute..."
cat > /app/.omniroute/accounts.json << ACCOUNTS
{"accounts": [{"name": "openai", "provider": "openai", "api_key": "${OPENAI_API_KEY:-dummy}", "default": true}]}
ACCOUNTS
omniroute start --port 20128 --accounts /app/.omniroute/accounts.json &
sleep 2

# Start Agent Zero Web UI
echo "🤖 Starting Agent Zero Web UI..."
cd /app && python3 run_ui.py &
sleep 5

# Start Marimo
echo "📓 Starting Marimo..."
marimo run /app/marimo/marimo_sandbox.py --host 0.0.0.0 --port 8081 --headless &
sleep 3

# Start proxy
echo "🔀 Starting proxy..."
python3 << 'PROXY'
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import httpx

app = FastAPI()
client = httpx.AsyncClient(timeout=120.0)

@app.get("/health")
async def health():
    return {"status": "healthy"}

@app.get("/marimo")
async def marimo_redir():
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/marimo/")

@app.api_route("/marimo/{path:path}", methods=["GET","POST","PUT","DELETE","PATCH"])
async def marimo_proxy(request: Request, path: str):
    url = f"http://localhost:8081/{path}"
    h = dict(request.headers)
    h.pop("host", None)
    b = await request.body() if request.method in ["POST","PUT","PATCH"] else None
    r = await client.request(request.method, url, headers=h, content=b)
    return StreamingResponse(r.aiter_bytes(), status_code=r.status_code, headers=dict(r.headers))

@app.api_route("/{path:path}", methods=["GET","POST","PUT","DELETE","PATCH"])
async def az_proxy(request: Request, path: str):
    url = f"http://localhost:80/{path}" if path else "http://localhost:80"
    h = dict(request.headers)
    h.pop("host", None)
    b = await request.body() if request.method in ["POST","PUT","PATCH"] else None
    r = await client.request(request.method, url, headers=h, content=b)
    return StreamingResponse(r.aiter_bytes(), status_code=r.status_code, headers=dict(r.headers))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)
PROXY

echo "✅ Ready!"
echo "   Agent Zero: http://localhost:7860"
echo "   Marimo: http://localhost:7860/marimo/"
wait
EOF

RUN chmod +x /app/start.sh

# 6. Environment
ENV HF_SPACE=true
ENV WORKSPACE=/app/workspace

# 7. Port
EXPOSE 7860

HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
    CMD curl -f http://localhost:7860/health || exit 1

ENTRYPOINT ["/app/start.sh"]