"""
Minimal API wrapper around CorvidaeAviary so the model does something visible
when deployed as an HF Docker Space.

Endpoints:
  GET  /             -- basic info
  GET  /health        -- readiness probe
  POST /forward       -- run a forward pass on random/user-supplied token ids
"""

import os
from typing import List, Optional

import torch
from fastapi import FastAPI
from pydantic import BaseModel

from corvid_aviary import CorvidaeAviary, SPECIALIST_ORDER

VOCAB_SIZE = 200
EMBEDDING_DIM = 64
MAX_SEQ_LEN = 32

app = FastAPI(title="Corvidae Aviary")

_model: Optional[CorvidaeAviary] = None


def get_model() -> CorvidaeAviary:
    global _model
    if _model is None:
        torch.manual_seed(0)
        _model = CorvidaeAviary(
            num_embeddings=VOCAB_SIZE, embedding_dim=EMBEDDING_DIM, max_seq_len=MAX_SEQ_LEN,
            nhead=4, dim_feedforward=128, memory_size=32, memory_word_size=16,
            num_read_heads=2, num_classes=VOCAB_SIZE, num_rook_experts=4, num_speakers=3,
            num_tracked_agents=3, raven_buffer_size=4, crow_num_tools=4, crow_max_steps=3,
            hippocampal_num_slots=16, hippocampal_coord_dim=3, statistical_memory_size=16,
            identity_capacity=8, num_surface_contexts=4,
        )
        _model.eval()
    return _model


class ForwardRequest(BaseModel):
    token_ids: Optional[List[int]] = None  # if omitted, random tokens are used
    seq_len: int = 16


@app.get("/")
def root():
    return {
        "name": "Corvidae Aviary",
        "specialists": SPECIALIST_ORDER,
        "endpoints": ["/health", "/forward"],
    }


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


@app.post("/forward")
def forward(req: ForwardRequest):
    model = get_model()
    if req.token_ids:
        ids = req.token_ids[: MAX_SEQ_LEN]
        x = torch.tensor([ids], dtype=torch.long)
    else:
        x = torch.randint(0, VOCAB_SIZE, (1, min(req.seq_len, MAX_SEQ_LEN)))

    model.reset_memory(batch_size=x.size(0))
    with torch.no_grad():
        logits, species_info = model(x, return_species_info=True)
        pred_ids = logits.argmax(dim=-1).squeeze(0).tolist()
        route_weights = species_info["route_weights"].mean(dim=(0, 1)).tolist()

    return {
        "input_ids": x.squeeze(0).tolist(),
        "predicted_next_ids": pred_ids,
        "route_weights_by_specialist": dict(zip(SPECIALIST_ORDER, route_weights)),
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))