#!/usr/bin/env python3
"""
Neuron Council Service — LLM anti-confabulation layer.

Fires 3 parallel Ollama calls and aggregates votes to produce a
confidence score + tags for any claim before it enters memory.

Usage:
    python3 council_service.py [--port 7771]
"""

import argparse
import asyncio
import time
import uuid
from typing import Optional

import httpx
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

OLLAMA_BASE = "http://localhost:11434/api/generate"
COUNCIL_MODELS = ["neuron:latest", "dolphin3:8b", "neuron-ft:latest"]
MODEL_TIMEOUT = 45.0  # seconds per model (models may need to load from cold)

SYSTEM_PROMPT_TEMPLATE = """\
You are a fact-checker. You will be given a claim.
Your job: assess if it is accurate, internally consistent, and grounded in reality.
Respond with EXACTLY ONE WORD:
- "plausible" if the claim seems accurate and well-grounded
- "uncertain" if you cannot determine accuracy or the claim is ambiguous
- "confabulation" if the claim appears to contain invented facts or clear errors

Claim: {claim}
Context: {context}

Your verdict (one word only):"""

VALID_VERDICTS = {"plausible", "uncertain", "confabulation"}

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Neuron Council Service",
    description="LLM-council anti-confabulation layer for Neuron soul",
    version="1.0.0",
)

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


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------

class VerifyRequest(BaseModel):
    claim: str
    context: Optional[str] = ""


class VerifyResponse(BaseModel):
    id: str
    claim: str
    confidence: float
    council_votes: list[str]
    summary: str
    tags: list[str]
    latency_ms: int


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------

async def query_model(client: httpx.AsyncClient, model: str, prompt: str) -> str:
    """
    Query a single Ollama model. Returns "plausible", "uncertain", or "confabulation".
    Returns "uncertain" on timeout. Raises httpx.ConnectError on connection failure.
    """
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
    }
    try:
        resp = await client.post(OLLAMA_BASE, json=payload, timeout=MODEL_TIMEOUT)
        resp.raise_for_status()
        data = resp.json()
        raw = data.get("response", "").strip().lower().split()[0] if data.get("response", "").strip() else "uncertain"
        # Normalise to one of the three valid verdicts
        if raw not in VALID_VERDICTS:
            return "uncertain"
        return raw
    except httpx.TimeoutException:
        return "uncertain"


async def run_council(claim: str, context: str) -> VerifyResponse:
    start = time.monotonic()
    prompt = SYSTEM_PROMPT_TEMPLATE.format(claim=claim, context=context)

    # Quick connectivity check — one tiny HEAD request to Ollama
    try:
        async with httpx.AsyncClient() as probe:
            await probe.get("http://localhost:11434/", timeout=2.0)
    except (httpx.ConnectError, httpx.TimeoutException):
        latency_ms = int((time.monotonic() - start) * 1000)
        return VerifyResponse(
            id=str(uuid.uuid4()),
            claim=claim,
            confidence=0.5,
            council_votes=[],
            summary="Ollama is unavailable; council could not convene.",
            tags=["council-unavailable"],
            latency_ms=latency_ms,
        )

    # Fire all 3 model calls in parallel
    async with httpx.AsyncClient() as client:
        tasks = [query_model(client, m, prompt) for m in COUNCIL_MODELS]
        votes: list[str] = await asyncio.gather(*tasks)

    plausible_count = votes.count("plausible")
    latency_ms = int((time.monotonic() - start) * 1000)

    # Voting rules
    if plausible_count == 3:
        confidence = 0.85
        tags = ["verified"]
        summary = "3/3 council members agree this is plausible."
    elif plausible_count == 2:
        confidence = 0.65
        tags = ["council-split"]
        summary = "2/3 council members agree this is plausible."
    elif plausible_count == 1:
        confidence = 0.30
        tags = ["unverified", "council-flagged"]
        summary = "1/3 council members found this plausible."
    else:
        confidence = 0.30
        tags = ["unverified", "council-flagged"]
        summary = "0/3 council members found this plausible."

    return VerifyResponse(
        id=str(uuid.uuid4()),
        claim=claim,
        confidence=confidence,
        council_votes=votes,
        summary=summary,
        tags=tags,
        latency_ms=latency_ms,
    )


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.post("/api/neuron/council/verify", response_model=VerifyResponse)
async def verify(req: VerifyRequest):
    return await run_council(req.claim, req.context or "")


@app.get("/healthz")
async def health():
    return {"status": "ok", "service": "council"}


# ---------------------------------------------------------------------------
# Startup warm-up: pre-load all council models so first real call is fast
# ---------------------------------------------------------------------------

@app.on_event("startup")
async def warmup_models():
    """
    Send a trivial prompt to each council model at startup.
    This forces Ollama to load the models into GPU memory so the first
    real council call does not pay the cold-load latency penalty.
    """
    print("[council] Warming up council models...")
    warmup_prompt = "Reply with one word: ready"
    async with httpx.AsyncClient() as client:
        tasks = [
            client.post(
                OLLAMA_BASE,
                json={"model": m, "prompt": warmup_prompt, "stream": False},
                timeout=60.0,
            )
            for m in COUNCIL_MODELS
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for model, result in zip(COUNCIL_MODELS, results):
            if isinstance(result, Exception):
                print(f"[council] warm-up failed for {model}: {result}")
            else:
                print(f"[council] {model} warm and ready")
    print("[council] All models warmed up.")


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Neuron Council Service")
    parser.add_argument("--port", type=int, default=7771, help="Port to listen on")
    parser.add_argument("--host", default="127.0.0.1", help="Host to bind to")
    args = parser.parse_args()

    print(f"[council] Starting on {args.host}:{args.port}")
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")