"""
app.py — the production FastAPI service (Knowledge Module 07 §1-3).

Run it:
    uvicorn app:app --reload
    curl -s localhost:8000/health
    curl -s -X POST localhost:8000/chat -H 'content-type: application/json' \
         -d '{"message":"How long to dispute a transaction?","conversation_id":"c1"}'
    curl -N -X POST localhost:8000/chat/stream -H 'content-type: application/json' \
         -d '{"message":"hello","conversation_id":"c2"}'

Demonstrates the production concerns the JD asks for:
  - async FastAPI + Pydantic contracts + auto OpenAPI at /docs
  - SSE token STREAMING (/chat/stream)
  - Cosmos-shaped conversation state (partition by conversation_id) — store.py
  - IDEMPOTENCY for safe retries (X-Idempotency-Key) — resilience.py
  - health/readiness probes for the orchestrator (Container Apps/AKS)

The responder here is a tiny grounded stand-in; in the capstone (Lab 07) the
LangGraph agent + RAG plug in behind the same endpoints.
"""
from __future__ import annotations

import asyncio
from typing import AsyncIterator

from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from resilience import IdempotencyCache
from store import CosmosLikeStore

app = FastAPI(title="Banking Assistant API", version="1.0")
store = CosmosLikeStore()                 # conversations, partitioned by conversation_id
idem = IdempotencyCache()
_seq: dict[str, int] = {}                 # per-conversation message counter


class ChatRequest(BaseModel):
    message: str = Field(min_length=1, max_length=4000)
    conversation_id: str = Field(min_length=1)
    language: str = "en"


class ChatResponse(BaseModel):
    conversation_id: str
    answer: str
    cited: bool


# A trivial grounded responder so the lab is self-contained. Lab 07 swaps in the
# real agent/RAG. It NEVER invents money facts — it answers from a tiny policy map
# or refers to a human.
_POLICY = {
    "dispute": "You can dispute a transaction within 60 days of the statement date. [dispute:0]",
    "fee": "The current account monthly fee is 25 AED, waived above 3000 AED. [fees:0]",
    "iban": "A UAE IBAN has 23 characters. [iban:0]",
}


def _respond(message: str) -> tuple[str, bool]:
    t = message.lower()
    for k, v in _POLICY.items():
        if k in t:
            return v, True
    return ("I don't have that in our documents; let me connect you with a human agent.",
            False)


def _next_seq(cid: str) -> int:
    _seq[cid] = _seq.get(cid, 0) + 1
    return _seq[cid]


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


@app.get("/ready")
async def ready() -> dict:
    # In prod, check downstream deps (AOAI/Search/Cosmos) reachability here.
    return {"ready": True, "store": store.stats()}


@app.post("/chat", response_model=ChatResponse)
async def chat(req: ChatRequest, x_idempotency_key: str | None = Header(default=None)) -> ChatResponse:
    # Idempotency: a retried request with the same key returns the cached reply
    # instead of writing a duplicate message (safe retries).
    if x_idempotency_key and (cached := idem.get(x_idempotency_key)):
        return ChatResponse(**cached)

    cid = req.conversation_id
    store.create_item({"id": f"{cid}:{_next_seq(cid)}", "conversationId": cid,
                       "role": "user", "content": req.message, "seq": _seq[cid]})
    answer, cited = _respond(req.message)
    store.create_item({"id": f"{cid}:{_next_seq(cid)}", "conversationId": cid,
                       "role": "assistant", "content": answer, "seq": _seq[cid]})

    resp = {"conversation_id": cid, "answer": answer, "cited": cited}
    if x_idempotency_key:
        idem.put(x_idempotency_key, resp)
    return ChatResponse(**resp)


@app.get("/conversations/{conversation_id}")
async def get_conversation(conversation_id: str) -> dict:
    # Single-partition point read — cheap because we partitioned by conversation_id.
    msgs = store.read_conversation(conversation_id)
    if not msgs:
        raise HTTPException(status_code=404, detail="conversation not found")
    return {"conversation_id": conversation_id, "messages": msgs,
            "ru_consumed": store.ru_consumed}


@app.post("/chat/stream")
async def chat_stream(req: ChatRequest) -> StreamingResponse:
    """SSE streaming: emit the answer token-by-token so the UI shows output
    immediately (perceived latency = time-to-first-token)."""
    answer, _ = _respond(req.message)

    async def gen() -> AsyncIterator[bytes]:
        for word in answer.split():
            await asyncio.sleep(0)            # yield to the event loop (real: model deltas)
            yield f"data: {word} \n\n".encode()
        yield b"data: [DONE]\n\n"

    return StreamingResponse(gen(), media_type="text/event-stream")