#!/usr/bin/env python3
"""
atlas_router_api.py  —  Atlas Chat backbone router  (v0.1  2026-06-02)
OpenAI-compatible /v1/chat/completions + /v1/models

Routes to:
  Anthropic  :  atlas-haiku  |  atlas-sonnet  |  atlas-opus
  Ollama     :  atlas-free (hermes3:8b)  |  atlas-code (qwen-coder)
              atlas-fast (llama3.2:1b)   |  atlas-qwen (qwen2.5:7b)
  Smart      :  atlas-auto → heuristic picks cheapest backend that fits

Run (local or atlas-01):
  python3 atlas_router_api.py --port 8888
  ATLAS_ROUTER_PORT=8888 python3 atlas_router_api.py

Env vars:
  ANTHROPIC_API_KEY   required for claude backends
  OLLAMA_BASE_URL     default http://localhost:11434
  ATLAS_ROUTER_PORT   default 8888
  ATLAS_HAIKU_MODEL   default claude-haiku-4-5-20251001
  ATLAS_SONNET_MODEL  default claude-sonnet-4-6
  ATLAS_OPUS_MODEL    default claude-opus-4-8
"""
from __future__ import annotations
import asyncio, json, logging, os, time, uuid
from pathlib import Path
import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Any
import uvicorn

# ── secrets (source /opt/atlas/router-secrets.env if present) ──────────────────
_sec = Path("/opt/atlas/router-secrets.env")
if not _sec.exists(): _sec = Path.home() / ".config/atlas/router-secrets.env"
if _sec.exists():
    for ln in _sec.read_text().splitlines():
        if "=" in ln and not ln.startswith("#"):
            k, v = ln.split("=", 1)
            os.environ[k.strip()] = v.strip()  # last value wins (normal dotenv behaviour)

ANTHROPIC_KEY  = os.environ.get("ANTHROPIC_API_KEY", "")
CLAUDE_RELAY   = os.environ.get("ATLAS_CLAUDE_RELAY", "")   # e.g. http://100.X.X.X:9191
OLLAMA_URL     = os.environ.get("OLLAMA_BASE_URL",  "http://localhost:11434")
PORT          = int(os.environ.get("ATLAS_ROUTER_PORT", "8888"))
LOG_LEVEL     = os.environ.get("ATLAS_LOG_LEVEL", "info")

M_HAIKU  = os.environ.get("ATLAS_HAIKU_MODEL",  "claude-haiku-4-5-20251001")
M_SONNET = os.environ.get("ATLAS_SONNET_MODEL", "claude-sonnet-4-6")
M_OPUS   = os.environ.get("ATLAS_OPUS_MODEL",   "claude-opus-4-8")

CLAUDE_MODELS = {
    "atlas-haiku":  M_HAIKU,
    "atlas-sonnet": M_SONNET,
    "atlas-opus":   M_OPUS,
}
OLLAMA_MODELS = {
    "atlas-free":  os.environ.get("ATLAS_FREE_MODEL", "qwen2.5:7b"),
    "atlas-code":  os.environ.get("ATLAS_CODE_MODEL", "qwen2.5-coder:7b"),
    "atlas-fast":  os.environ.get("ATLAS_FAST_MODEL", "qwen2.5:1.5b"),
    "atlas-qwen":  os.environ.get("ATLAS_QWEN_MODEL", "qwen2.5:7b"),
    "atlas-phi3":  os.environ.get("ATLAS_PHI3_MODEL", "qwen2.5:7b"),
}
ALL_ATLAS = list(CLAUDE_MODELS) + list(OLLAMA_MODELS) + ["atlas-auto"]

logging.basicConfig(level=LOG_LEVEL.upper(), format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("atlas-router")
app = FastAPI(title="Atlas Router", version="0.1")

# ── Pydantic models ──────────────────────────────────────────────────────────────
class Msg(BaseModel):
    role: str
    content: Any  # str or list[dict] for multimodal

class ChatRequest(BaseModel):
    model: str = "atlas-auto"
    messages: list[Msg]
    max_tokens: int = 4096
    temperature: float = 0.7
    stream: bool = True
    system: str | None = None

# ── Auto-router heuristic ────────────────────────────────────────────────────────
def _auto_route(messages: list[Msg]) -> str:
    text = " ".join(
        (m.content if isinstance(m.content, str) else json.dumps(m.content))
        for m in messages
    ).lower()
    wc = len(text.split())

    CODE  = {"```","def ","function","class ","import ","bug","implement","refactor",
             "sql","python","javascript","typescript","bash","powershell","html","css"}
    PLAN  = {"plan","design","architect","strategy","roadmap","analyse","analyze",
             "overview","approach","decide","compare","when will","how should"}

    has_code  = any(s in text for s in CODE)
    needs_plan = (any(s in text for s in PLAN) or wc > 200)

    if wc < 12 and not has_code:
        return "atlas-fast"    # ultra-short → free 1B
    if has_code and not needs_plan:
        return "atlas-code"    # code → free qwen-coder
    if needs_plan and wc > 120:
        return "atlas-sonnet"  # heavy planning → sonnet
    if needs_plan:
        return "atlas-haiku"   # light analysis → haiku (cheap)
    if not ANTHROPIC_KEY:
        return "atlas-free"    # no key → always free
    return "atlas-haiku"       # default: cheap + fast

# ── SSE helpers ──────────────────────────────────────────────────────────────────
def _chunk(cid: str, model: str, content: str, finish: str | None = None) -> str:
    return "data: " + json.dumps({
        "id": cid, "object": "chat.completion.chunk",
        "created": int(time.time()), "model": model,
        "choices": [{"index": 0,
                     "delta": {"content": content} if content else {},
                     "finish_reason": finish}]
    }) + "\n\n"

def _done(): return "data: [DONE]\n\n"

# ── Anthropic streaming ──────────────────────────────────────────────────────────
async def _stream_claude(req: ChatRequest, claude_model: str):
    if not ANTHROPIC_KEY:
        yield _chunk("err", req.model, "[atlas-router] ANTHROPIC_API_KEY not set — switch to atlas-free")
        yield _done(); return
    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    # build message list (strip system from messages if present)
    sys_prompt = req.system or next(
        (m.content for m in req.messages if m.role == "system" and isinstance(m.content,str)), None)
    msgs = [{"role": m.role, "content": m.content}
            for m in req.messages if m.role != "system"]
    headers = {
        "x-api-key": ANTHROPIC_KEY,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }
    body: dict[str, Any] = {
        "model": claude_model,
        "max_tokens": req.max_tokens,
        "temperature": req.temperature,
        "messages": msgs,
        "stream": True,
    }
    if sys_prompt: body["system"] = sys_prompt
    async with httpx.AsyncClient(timeout=120) as client:
        async with client.stream("POST", "https://api.anthropic.com/v1/messages",
                                 headers=headers, json=body) as resp:
            if resp.status_code != 200:
                err = await resp.aread()
                log.error("anthropic %s: %s", resp.status_code, err[:200])
                yield _chunk(cid, req.model, f"[atlas-router] Anthropic error {resp.status_code}")
                yield _done(); return
            async for line in resp.aiter_lines():
                if not line or not line.startswith("data:"): continue
                raw = line[5:].strip()
                if raw == "[DONE]": break
                try:
                    ev = json.loads(raw)
                except Exception: continue
                ev_type = ev.get("type","")
                if ev_type == "content_block_delta":
                    delta = ev.get("delta", {})
                    if delta.get("type") == "text_delta":
                        yield _chunk(cid, req.model, delta.get("text", ""))
                elif ev_type == "message_stop":
                    break
    yield _chunk(cid, req.model, "", "stop")
    yield _done()
    log.info("claude/%s done cid=%s", claude_model, cid)

# ── Model availability cache ─────────────────────────────────────────────────────
_avail_cache: dict[str, bool] = {}
_avail_ts: float = 0

def _ollama_model_available(model_name: str) -> bool:
    """Check if a model exists on the Ollama instance (cached for 60s)."""
    global _avail_ts
    if time.time() - _avail_ts < 60 and model_name in _avail_cache:
        return _avail_cache[model_name]
    try:
        import urllib.request
        req = urllib.request.Request(f"{OLLAMA_URL}/api/tags")
        with urllib.request.urlopen(req, timeout=5) as r:
            tags = json.loads(r.read())
            available = {m["name"] for m in tags.get("models", [])}
            for m in available:
                _avail_cache[m] = True
            _avail_ts = time.time()
            return model_name in available
    except Exception:
        return True  # assume available if we can't check

# ── Circuit breaker ───────────────────────────────────────────────────────────────
_failures: dict[str, list[float]] = {}

def _breaker_trip(backend: str) -> bool:
    """Return True if the backend has 3+ failures in 5 minutes (should skip)."""
    now = time.time()
    _failures.setdefault(backend, [])
    _failures[backend] = [t for t in _failures[backend] if now - t < 300]
    return len(_failures[backend]) >= 3

def _breaker_record(backend: str):
    """Record a failure for the given backend."""
    _failures.setdefault(backend, []).append(time.time())

def _breaker_reset(backend: str):
    """Reset failures on success."""
    _failures.pop(backend, None)

# ── Ollama streaming ─────────────────────────────────────────────────────────────
async def _stream_ollama(req: ChatRequest, ollama_model: str):
    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    msgs = [{"role": m.role, "content": m.content if isinstance(m.content,str)
             else json.dumps(m.content)} for m in req.messages]
    body = {"model": ollama_model, "messages": msgs, "stream": True,
            "options": {"temperature": req.temperature, "num_predict": req.max_tokens}}
    breaker_key = f"ollama-{ollama_model}"
    if _breaker_trip(breaker_key):
        yield _chunk(cid, req.model, f"[atlas-router] Circuit open for {ollama_model} — too many failures")
        yield _done(); return
    try:
        async with httpx.AsyncClient(timeout=120) as client:
            async with client.stream("POST", f"{OLLAMA_URL}/api/chat", json=body) as resp:
                if resp.status_code != 200:
                    err = await resp.aread()
                    log.error("ollama %s: %s", resp.status_code, err[:200])
                    _breaker_record(breaker_key)
                    yield _chunk(cid, req.model, f"[atlas-router] Ollama error {resp.status_code}")
                    yield _done(); return
                async for line in resp.aiter_lines():
                    if not line: continue
                    try: ev = json.loads(line)
                    except Exception: continue
                    content = (ev.get("message") or {}).get("content", "")
                    if content: yield _chunk(cid, req.model, content)
                    if ev.get("done"): break
    except Exception as e:
        log.error("ollama %s: %s", ollama_model, e)
        _breaker_record(breaker_key)
        yield _chunk(cid, req.model, f"[atlas-router] Ollama error {e}")
        yield _done(); return
    _breaker_reset(breaker_key)
    yield _chunk(cid, req.model, "", "stop")
    yield _done()
    log.info("ollama/%s done cid=%s", ollama_model, cid)

# ── Main dispatch ────────────────────────────────────────────────────────────────
async def _stream_relay(req: ChatRequest, relay_url: str):
    """Forward Claude-tier request to the local Claude CLI relay (Max subscription)."""
    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    body = {"model": req.model, "messages": [{"role": m.role, "content": m.content} for m in req.messages],
            "max_tokens": req.max_tokens, "stream": False}
    try:
        async with httpx.AsyncClient(timeout=130) as client:
            resp = await client.post(f"{relay_url}/v1/chat/completions", json=body)
            data = resp.json()
            content = data["choices"][0]["message"]["content"]
    except Exception as e:
        content = f"[atlas-relay unreachable: {e}] — is Tailscale connected and relay running?"
    yield _chunk(cid, req.model, content)
    yield _chunk(cid, req.model, "", "stop")
    yield _done()

def _dispatch(req: ChatRequest):
    model = req.model
    if model == "atlas-auto":
        model = _auto_route(req.messages)
        log.info("auto-route → %s  (wc=%d)", model, sum(len((m.content or "").split())
                 for m in req.messages if isinstance(m.content, str)))
    if model in CLAUDE_MODELS:
        if CLAUDE_RELAY:
            log.info("routing %s → relay %s", model, CLAUDE_RELAY)
            return _stream_relay(req, CLAUDE_RELAY)
        elif ANTHROPIC_KEY:
            return _stream_claude(req, CLAUDE_MODELS[model])
        else:
            # Fallback: downgrade to atlas-free with a note
            log.warning("%s: no relay or API key — downgrading to atlas-free", model)
            req_copy = req.model_copy(update={"model": "atlas-free"})
            return _stream_ollama(req_copy, OLLAMA_MODELS["atlas-free"])
    if model in OLLAMA_MODELS:
        ollama_model = OLLAMA_MODELS[model]
        if not _ollama_model_available(ollama_model) and not _breaker_trip(f"ollama-{ollama_model}"):
            log.warning("%s: model %s not available on ollama — pulling...", model, ollama_model)
            try:
                import urllib.request as _ur
                _ur.urlopen(_ur.Request(f"{OLLAMA_URL}/api/pull",
                    data=json.dumps({"name": ollama_model}).encode(),
                    headers={"Content-Type": "application/json"}, method="POST"), timeout=120)
                _avail_ts = 0  # force cache refresh
            except Exception as e:
                log.error("pull failed: %s", e)
            if not _ollama_model_available(ollama_model):
                fallback = next((m for m in ["qwen2.5:1.5b","llama3.2:3b"] if _ollama_model_available(m)), ollama_model)
                log.warning("falling back to %s", fallback)
                return _stream_ollama(req, fallback)
        return _stream_ollama(req, ollama_model)
    # fallback: treat model string as raw ollama model name
    log.warning("unknown atlas model %r — forwarding as ollama", model)
    return _stream_ollama(req, model)

# ── Endpoints ────────────────────────────────────────────────────────────────────
@app.get("/v1/models")
async def list_models():
    models = []
    ts = int(time.time())
    labels = {
        "atlas-auto":   "Atlas Auto (smart router)",
        "atlas-haiku":  "Atlas Fast (Haiku)",
        "atlas-sonnet": "Atlas Pro (Sonnet)",
        "atlas-opus":   "Atlas Max (Opus)",
        "atlas-free":   "Atlas Free (hermes3 local)",
        "atlas-code":   "Atlas Code (Qwen Coder local)",
        "atlas-fast":   "Atlas Ultra-Fast (LLaMA 1B local)",
        "atlas-qwen":   "Atlas Qwen (qwen2.5 local)",
        "atlas-phi3":   "Atlas Reason (qwen2.5:7b local)",
    }
    for mid in ALL_ATLAS:
        models.append({"id": mid, "object": "model", "created": ts,
                       "owned_by": "atlas-corporation",
                       "description": labels.get(mid, mid)})
    return {"object": "list", "data": models}

@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    log.info("→ model=%s msgs=%d", req.model, len(req.messages))
    gen = _dispatch(req)
    if req.stream:
        return StreamingResponse(gen, media_type="text/event-stream",
                                 headers={"Cache-Control": "no-cache",
                                          "X-Accel-Buffering": "no"})
    # non-streaming: collect full response
    content = ""
    async for chunk in gen:
        if chunk.startswith("data:") and not chunk.strip().endswith("[DONE]"):
            try:
                ev = json.loads(chunk[5:].strip())
                content += (ev["choices"][0]["delta"] or {}).get("content", "")
            except Exception: pass
    return {"id": f"chatcmpl-{uuid.uuid4().hex[:12]}", "object": "chat.completion",
            "created": int(time.time()), "model": req.model,
            "choices": [{"index": 0, "message": {"role": "assistant", "content": content},
                         "finish_reason": "stop"}],
            "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1}}

@app.get("/health")
async def health():
    return {"status": "ok", "models": len(ALL_ATLAS),
            "anthropic": bool(ANTHROPIC_KEY), "ollama": OLLAMA_URL}

@app.get("/healthz")
async def healthz():
    """Full system health: disk, RAM, swap, Docker, uptime."""
    import shutil, subprocess
    disk = shutil.disk_usage("/")
    mem = {}
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if ":" in line:
                    k, v = line.split(":", 1)
                    mem[k.strip()] = int(v.replace("kB","").strip()) // 1024
    except: pass
    docker_ok, docker_total, docker_running = False, 0, 0
    try:
        r = subprocess.run(["docker","ps","-a","--format","{{.Status}}"],
                          capture_output=True, text=True, timeout=5)
        containers = r.stdout.strip().split("\n") if r.stdout.strip() else []
        docker_total = len(containers)
        docker_running = sum(1 for c in containers if c.startswith("Up"))
        docker_ok = True
    except: pass
    ollama_models = 0
    try:
        import urllib.request as _ur
        r = _ur.urlopen(f"{OLLAMA_URL}/api/tags", timeout=3)
        ollama_models = len(json.loads(r.read()).get("models", []))
    except: pass
    uptime_s = time.time() - _start_time if "_start_time" in dir() else 0
    return {
        "status": "ok" if docker_ok and ollama_models > 0 else "degraded",
        "disk_pct": disk.used / disk.total * 100 if disk.total else 0,
        "disk_free_gb": disk.free / 1e9,
        "disk_total_gb": disk.total / 1e9,
        "ram_used_mb": mem.get("MemTotal", 0) - mem.get("MemAvailable", 0),
        "ram_total_mb": mem.get("MemTotal", 0),
        "swap_used_mb": mem.get("SwapTotal", 0) - mem.get("SwapFree", 0),
        "swap_total_mb": mem.get("SwapTotal", 0),
        "docker_containers_running": docker_running,
        "docker_containers_total": docker_total,
        "ollama_models": ollama_models,
        "uptime_s": uptime_s,
        "router_models": len(ALL_ATLAS),
    }

_start_time = time.time()

# ── Entry point ───────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=PORT)
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--reload", action="store_true")
    a = ap.parse_args()
    log.info("Atlas Router v0.1  port=%d  anthropic=%s  ollama=%s",
             a.port, bool(ANTHROPIC_KEY), OLLAMA_URL)
    uvicorn.run("atlas_router_api:app" if a.reload else app,
                host=a.host, port=a.port, reload=a.reload,
                log_level=LOG_LEVEL, access_log=True)