main.py

"""
Instant-RAG Platform — Agent-first RAG with receipts, swarm reasoning & deterministic cost.
"""
from fastapi import FastAPI, UploadFile, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional
import logging
import os
import time

from tenants.manager import tm
from core.retriever import SimpleRetriever, chunk_text, weave_answer
from core.ratelimit import limiter
from core.subscription import subs
from core.audit import auditor
from contracts.engine import engine
from explain.trace import build_trace
from explain.scores import confidence_from_parts
from ethics.judge import judge
from identity.passport import passport
from swarm.session import run as swarm_run
from api_trust import router as trust_router

# ─── Logging ─────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# ─── App ─────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Instant-RAG Platform",
    description=(
        "Agent-first RAG with Polygon USDC micropayments, swarm reasoning, "
        "and verifiable receipts. Upload docs, query with citations, pay per call."
    ),
    version="1.1.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Routers (register BEFORE static mounts — order is critical) ─────────
# app.mount("/") is a catch-all. Any route added after it will be shadowed.
# Always include_router first, mount static last.
app.include_router(trust_router)
from dashboard import router as admin_router
app.include_router(admin_router)

# ─── Helpers ─────────────────────────────────────────────────────────────

def _verify_or_401(agent_id: str, token: str):
    if not passport.verify(agent_id, token):
        raise HTTPException(status_code=401, detail="invalid_passport")

def _check_subscription(agent_id: str):
    if subs.check(agent_id) != "active":
        raise HTTPException(status_code=403, detail="subscription_inactive")

# ─── Models ──────────────────────────────────────────────────────────────

class QueryRequest(BaseModel):
    text: str = Field(..., max_length=10000, min_length=1)
    agent_id: str = Field(..., min_length=1, max_length=100)
    token: str = Field(..., min_length=1)

class RegisterRequest(BaseModel):
    agent_id: str = Field(..., min_length=1, max_length=100)
    role: Optional[str] = Field(default="agent")

class SpendRequest(BaseModel):
    agent_id: str
    action: str

# ─── System ───────────────────────────────────────────────────────────────

@app.get("/health", tags=["system"])
def health():
    return {"status": "ok", "version": "1.1.0", "timestamp": time.time()}

@app.get("/ready", tags=["system"])
async def ready():
    return {"status": "ok"}

# ─── Identity ────────────────────────────────────────────────────────────

@app.post("/register", tags=["identity"], summary="Register agent, get auth token")
def register(req: RegisterRequest):
    """
    Register a new agent. Returns a token — store it securely, shown only once.
    """
    try:
        token = passport.issue(req.agent_id, role=req.role or "agent")
        subs.activate(req.agent_id)
        auditor.record("register", req.agent_id, {"role": req.role})
        return {
            "agent_id": req.agent_id,
            "token": token,
            "warning": "Store this token securely. It is not recoverable.",
            "next_steps": {
                "ingest": "POST /ingest — upload a document",
                "query": "POST /query — ask a question"
            }
        }
    except Exception as e:
        logger.error(f"Registration error: {e}")
        raise HTTPException(status_code=500, detail="registration_failed")

# ─── Core RAG ────────────────────────────────────────────────────────────

@app.post("/ingest", tags=["rag"], summary="Index a UTF-8 document into agent memory")
async def ingest(file: UploadFile, agent_id: str, token: str):
    try:
        _verify_or_401(agent_id, token)
        _check_subscription(agent_id)

        MAX_BYTES = 10 * 1024 * 1024  # 10 MB cap (Render free tier friendly)
        content = await file.read(MAX_BYTES + 1)
        if len(content) > MAX_BYTES:
            raise HTTPException(status_code=413, detail="file_too_large_10mb_limit")

        try:
            text = content.decode("utf-8")
        except UnicodeDecodeError:
            raise HTTPException(status_code=400, detail="file_must_be_utf8_text")

        if not text.strip():
            raise HTTPException(status_code=400, detail="file_is_empty")

        chunks = chunk_text(text)
        tenant = tm.get(agent_id)
        tenant.retriever.add_documents(chunks, source_name=file.filename or "unknown")

        auditor.record("ingest", agent_id, {
            "chunks": len(chunks),
            "filename": file.filename,
            "size": len(text)
        })

        return {
            "status": "indexed",
            "chunks": len(chunks),
            "filename": file.filename,
            "agent_id": agent_id
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Ingestion error: {e}")
        raise HTTPException(status_code=500, detail="ingestion_failed")


@app.post("/query", tags=["rag"], summary="Semantic search over agent memory")
async def query(request: QueryRequest):
    try:
        _verify_or_401(request.agent_id, request.token)
        _check_subscription(request.agent_id)

        ok, reason = judge.inspect(request.text)
        if not ok:
            auditor.record("ethics_block", request.agent_id, {
                "query": request.text[:120], "reason": reason
            })
            raise HTTPException(status_code=400, detail=f"ethics_block: {reason}")

        if not limiter.allow(request.agent_id):
            usage = limiter.get_usage(request.agent_id)
            raise HTTPException(
                status_code=429,
                detail=f"rate_limited: {usage['used']}/{usage['limit']} daily",
                headers={"Retry-After": "3600"}
            )

        tenant = tm.get(request.agent_id)

        if not tenant.retriever.docs:
            return {
                "answer": "No documents indexed yet. Use POST /ingest to add documents first.",
                "citations": [],
                "confidence": 0.0,
                "explanation": {"method": "none", "steps": []}
            }

        results, cites, scores = tenant.retriever.search(request.text)
        trace = build_trace(request.text, results, scores)
        packet = weave_answer(results, cites)
        packet["explanation"] = trace
        packet["confidence"] = confidence_from_parts(
            0.7, max(scores) if scores else 0, len(cites)
        )

        auditor.record("query", request.agent_id, {
            "q": request.text[:120],
            "results": len(results),
            "confidence": packet["confidence"]
        })

        return packet

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Query error: {e}")
        raise HTTPException(status_code=500, detail="query_failed")


@app.post("/swarm/query", tags=["swarm"], summary="Parallel swarm reasoning query")
async def swarm_query(request: QueryRequest):
    try:
        _verify_or_401(request.agent_id, request.token)
        _check_subscription(request.agent_id)

        async def query_fn(text: str):
            q = QueryRequest(text=text, agent_id=request.agent_id, token=request.token)
            return await query(q)

        result = await swarm_run(request.text, query_fn)
        auditor.record("swarm_query", request.agent_id, {"q": request.text[:120]})
        return result

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Swarm error: {e}")
        raise HTTPException(status_code=500, detail="swarm_query_failed")

# ─── Analytics ────────────────────────────────────────────────────────────

@app.get("/stats/{agent_id}", tags=["analytics"], summary="Agent usage stats")
async def get_stats(agent_id: str, token: str):
    try:
        _verify_or_401(agent_id, token)
        tenant = tm.get(agent_id)
        agent_logs = auditor.read_by_agent(agent_id)
        usage = limiter.get_usage(agent_id)

        return {
            "agent_id": agent_id,
            "total_queries": len([l for l in agent_logs if l.get("event") == "query"]),
            "total_swarm_queries": len([l for l in agent_logs if l.get("event") == "swarm_query"]),
            "total_ingestions": len([l for l in agent_logs if l.get("event") == "ingest"]),
            "total_documents": len(tenant.retriever.docs),
            "rate_limit": usage,
            "subscription_plan": subs.get_plan(agent_id)
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Stats error: {e}")
        raise HTTPException(status_code=500, detail="stats_failed")

# ─── Payments ─────────────────────────────────────────────────────────────

from payments.ledger import balance, spend
from payments.pricing import prices

@app.get("/wallet/balance", tags=["payments"])
def get_balance(agent_id: str, token: str):
    _verify_or_401(agent_id, token)
    return {"agent_id": agent_id, "balance_usdc": balance(agent_id)}

@app.post("/wallet/spend", tags=["payments"])
def wallet_spend(req: SpendRequest):
    cost = prices.get(req.action)
    if cost is None:
        return {"error": "unknown_action", "valid_actions": list(prices.keys())}
    if not spend(req.agent_id, cost):
        return {"error": "insufficient_funds", "required": cost, "balance": balance(req.agent_id)}
    return {
        "status": "ok",
        "action": req.action,
        "cost_usdc": cost,
        "remaining_usdc": balance(req.agent_id),
        "timestamp": time.time()
    }

# ─── Static / Agent Discovery (MUST be last — / is a catch-all) ─────────
from fastapi.staticfiles import StaticFiles

os.makedirs("static/.well-known", exist_ok=True)

app.mount("/.well-known", StaticFiles(directory="static/.well-known"), name="wellknown")
app.mount("/", StaticFiles(directory="static", html=True), name="public")

# ─── Run ──────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)