"""RAG pipeline example: inject many context blocks, let the compiler handle truncation.

This simulates a retrieval-augmented generation flow where you have more context
chunks than will fit in the budget. The compiler drops lowest-relevance blocks first.
"""

from prompt_builder import (
    Prompt,
    Persona,
    Instruction,
    InputData,
    OutputSpec,
    ContextBlock,
    TokenBudget,
    compile,
)


def simulate_rag_retrieval(query: str) -> list[ContextBlock]:
    """Pretend we got these chunks from a vector DB search."""
    return [
        ContextBlock(
            content="Sockudo uses a sharded architecture with consistent hashing for channel distribution across nodes.",
            source="docs/architecture.md",
            relevance=0.95,
            metadata={"section": "sharding"},
        ),
        ContextBlock(
            content="WebSocket connections are upgraded from HTTP/1.1. The server must respond with 101 Switching Protocols.",
            source="docs/protocol.md",
            relevance=0.88,
            metadata={"section": "handshake"},
        ),
        ContextBlock(
            content="Rate limiting uses a token bucket algorithm with configurable burst size per channel.",
            source="docs/rate-limiting.md",
            relevance=0.72,
            metadata={"section": "rate-limits"},
        ),
        ContextBlock(
            content="The project started in 2024 as a Pusher-compatible alternative written in Rust.",
            source="docs/history.md",
            relevance=0.30,
            metadata={"section": "background"},
        ),
        ContextBlock(
            content="CI pipeline runs on GitHub Actions with cross-compilation for linux-amd64 and linux-arm64.",
            source="docs/contributing.md",
            relevance=0.15,
            metadata={"section": "ci"},
        ),
        ContextBlock(
            content="Benchmarks show 50k concurrent connections on a single node with p99 latency under 5ms.",
            source="docs/benchmarks.md",
            relevance=0.82,
            metadata={"section": "performance"},
        ),
    ]


def answer_question(question: str):
    """Build a RAG prompt and compile it."""

    chunks = simulate_rag_retrieval(question)

    prompt = Prompt(
        persona=Persona(
            role="Technical documentation assistant for the Sockudo project.",
            tone="precise, helpful",
        ),
        instruction=Instruction(
            task="Answer the user's question based on the provided documentation context.",
            constraints=[
                "Only use information from the provided context blocks",
                "If the context doesn't contain enough info, say so",
            ],
        ),
        input_data=InputData(content=question, label="question"),
        output_spec=OutputSpec(format="markdown", length="2-3 paragraphs"),
        context=chunks,
    )

    # Generous budget — all 6 chunks fit
    result_full = compile(prompt, TokenBudget(context_window=8_000))
    print("=== Generous budget (8k tokens) ===")
    print(result_full.report())

    # Tight budget — compiler must drop some chunks
    result_tight = compile(prompt, TokenBudget(context_window=1_000, output_reserve=300))
    print("\n=== Tight budget (1k tokens) ===")
    print(result_tight.report())

    # See exactly what survived
    print("\nWarnings:")
    for w in result_tight.warnings:
        print(f"  → {w}")

    print("\nFinal user message:")
    user_msg = [m for m in result_tight.messages if m["role"] == "user"][0]
    print(user_msg["content"])


if __name__ == "__main__":
    answer_question("How does Sockudo handle scaling across multiple nodes?")
