from pydantic import BaseModel, Field


class Citation(BaseModel):
    chunk_id: str = Field(
        description="ID of the source chunk (must match an ID from the provided context)"
    )
    excerpt: str = Field(
        description="Verbatim quote (≤150 chars) from the chunk that directly supports the answer"
    )
    relevance_score: float = Field(
        ge=0.0, le=1.0,
        description="How essential this chunk is to the answer (0=tangential, 1=directly answers it)",
    )


class StructuredAnswer(BaseModel):
    answer: str = Field(
        description="Direct, concise answer to the question in 1-3 sentences"
    )
    citations: list[Citation] = Field(
        description="Chunks that were used to construct the answer - omit irrelevant chunks"
    )
    confidence: float = Field(
        ge=0.0, le=1.0,
        description="How fully the provided context supports the answer (0=guessing, 1=fully supported)",
    )
    cannot_answer: bool = Field(
        description="Set true when the context lacks sufficient information - do NOT hallucinate"
    )
    reasoning: str = Field(
        description="One sentence explaining how the answer was derived from the cited chunks"
    )

import json
from schemas import StructuredAnswer

# The system prompt template - {schema} is replaced at call time
_SYSTEM_TEMPLATE = """\
You are a precise question-answering assistant. Answer the question using ONLY the provided context.

You MUST respond with valid JSON that exactly matches this schema:
{schema}

Rules:
- Set cannot_answer=true if the context lacks enough information to answer confidently
- citations must reference real chunk_ids from the context - never invent IDs
- excerpt must be a verbatim quote (≤150 chars), not a paraphrase
- confidence reflects context coverage, not your general knowledge
- Output ONLY the JSON object - no explanation, no markdown fences"""

# At query time:
schema_str = json.dumps(StructuredAnswer.model_json_schema(), indent=2)
system     = _SYSTEM_TEMPLATE.format(schema=schema_str)

{
  "type": "object",
  "properties": {
    "answer": {
      "type": "string",
      "description": "Direct, concise answer to the question in 1-3 sentences"
    },
    "confidence": {
      "type": "number",
      "minimum": 0.0,
      "maximum": 1.0,
      "description": "How fully the provided context supports the answer (0=guessing, 1=fully supported)"
    },
    "cannot_answer": {
      "type": "boolean",
      "description": "Set true when the context lacks sufficient information - do NOT hallucinate"
    },
    ...
  },
  "required": ["answer", "citations", "confidence", "cannot_answer", "reasoning"]
}

def query(
    question:   str,
    collection: chromadb.Collection,
    model:      str = "llama3.2:3b",
    k:          int = 5,
) -> StructuredAnswer:
    # 1. Retrieve k chunks
    chunks = _vector_search(collection, question, k)
    if not chunks:
        return StructuredAnswer(
            answer="No documents found in this collection.",
            citations=[], confidence=0.0, cannot_answer=True,
            reasoning="vector search returned no results",
        )

    # 2. Build schema-aware system prompt
    schema_str = json.dumps(StructuredAnswer.model_json_schema(), indent=2)
    context    = "\n\n".join(f"[{cid}]\n{text}" for cid, text in chunks)
    system     = _SYSTEM_TEMPLATE.format(schema=schema_str)

    # 3. Call Ollama with format="json" (grammar-constrained decoding)
    resp = ollama.chat(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {question}"},
        ],
        format="json",
        options={"temperature": 0, "num_predict": 1024},
    )
    raw = resp["message"]["content"].strip()

    # 4. Validate with Pydantic
    return _parse(raw)

def _parse(raw: str) -> StructuredAnswer:
    # Layer 1 - happy path: full Pydantic validation
    try:
        return StructuredAnswer.model_validate_json(raw)
    except (ValidationError, json.JSONDecodeError) as exc:
        log.warning("Structured parse failed (%s): %.120s", type(exc).__name__, raw)

        # Layer 2 - partial recovery: salvage the answer field if present
        try:
            data = json.loads(raw)
            return StructuredAnswer(
                answer=str(data.get("answer", raw[:400])),
                citations=[],
                confidence=0.0,
                cannot_answer=True,
                reasoning="schema validation failed - partial response preserved",
            )
        except Exception:
            # Layer 3 - last resort: raw text as answer
            return StructuredAnswer(
                answer=raw[:400],
                citations=[], confidence=0.0,
                cannot_answer=True, reasoning="json parse error",
            )

router = APIRouter(prefix="/structured", tags=["structured"])

@router.post("/query", response_model=QueryResponse)
async def structured_query(
    q:         Annotated[str, Query(description="Question to answer")],
    tenant_id: Annotated[str, Query()] = "default",
    model:     Annotated[str, Query(description="Ollama model")] = "llama3.2:3b",
    k:         Annotated[int, Query(ge=1, le=10)] = 5,
):
    collection = _get_collection(tenant_id)
    result     = sr.query(question=q, collection=collection, model=model, k=k)
    return QueryResponse(
        answer        = result.answer,
        citations     = [CitationOut(...) for c in result.citations],
        confidence    = round(result.confidence, 3),
        cannot_answer = result.cannot_answer,
        reasoning     = result.reasoning,
    )


@router.get("/schema")
async def get_answer_schema() -> dict:
    # Expose the schema - useful for API consumers and debugging prompts
    return StructuredAnswer.model_json_schema()

# Run a structured query
curl -s -X POST \
  "http://localhost:8000/structured/query?q=What+is+the+rate+limit+policy&tenant_id=default&k=5" \
  | python3 -m json.tool

# Inspect the live schema the LLM receives
curl -s "http://localhost:8000/structured/schema" | python3 -m json.tool