"""Custom priority map + compilation inspection.

Shows how to override the default truncation order and use the
CompilationResult to understand exactly what the compiler did.
"""

from prompt_builder import (
    Prompt,
    Persona,
    Instruction,
    InputData,
    OutputSpec,
    ContextBlock,
    Example,
    Message,
    TokenBudget,
    PriorityMap,
    PriorityEntry,
    Section,
    TruncationStrategy,
    compile,
)


def build_chatbot_prompt(user_message: str, chat_history: list[dict]):
    """A chatbot where history matters more than RAG context.

    Default priority cuts context first, then history. Here we flip that:
    we'd rather keep the full conversation and drop some docs.
    But we also protect examples (never cut) since our few-shot
    demonstrations are critical for output quality.
    """

    prompt = Prompt(
        persona=Persona(
            role="Helpful customer support agent for a SaaS product.",
            tone="friendly, empathetic",
        ),
        instruction=Instruction(
            task="Respond to the customer's latest message.",
            constraints=["Don't make promises about features or timelines"],
        ),
        output_spec=OutputSpec(format="plain", length="2-4 sentences"),

        # Knowledge from help docs
        context=[
            ContextBlock(
                content="Billing cycles reset on the 1st of each month. Prorated charges apply for mid-cycle upgrades.",
                source="help/billing.md",
                relevance=0.7,
            ),
            ContextBlock(
                content="To cancel, go to Settings > Subscription > Cancel Plan. Cancellations take effect at end of billing period.",
                source="help/cancellation.md",
                relevance=0.9,
            ),
            ContextBlock(
                content="Enterprise plans include SSO, audit logs, and dedicated support. Contact sales@example.com.",
                source="help/enterprise.md",
                relevance=0.3,
            ),
        ],

        # Few-shot examples — we want these protected
        examples=[
            Example(
                input="I want to cancel my account",
                output="I'm sorry to hear that! You can cancel anytime from Settings > Subscription. Your access continues until the end of your current billing period. Is there anything I can help with before you go?",
                weight=1.0,
            ),
        ],

        # Conversation history
        history=[Message(role=m["role"], content=m["content"]) for m in chat_history],

        # Input
        input_data=InputData(content=user_message, label="customer_message"),

        # CUSTOM PRIORITY: protect history and examples, sacrifice context first
        priority=PriorityMap(entries=[
            PriorityEntry(section=Section.CONTEXT, strategy=TruncationStrategy.CUT_BY_RELEVANCE),
            # ↑ cut docs first (lowest relevance dropped)
            PriorityEntry(section=Section.HISTORY, strategy=TruncationStrategy.TRUNCATE_OLDEST),
            # ↑ only if context is fully exhausted, then trim old messages
            PriorityEntry(section=Section.EXAMPLES, strategy=TruncationStrategy.NEVER_CUT),
            # ↑ examples are PROTECTED (override default which allows cutting)
            PriorityEntry(section=Section.COT, strategy=TruncationStrategy.CUT_LAST),
            PriorityEntry(section=Section.PERSONA, strategy=TruncationStrategy.NEVER_CUT),
            PriorityEntry(section=Section.RULES, strategy=TruncationStrategy.NEVER_CUT),
            PriorityEntry(section=Section.INSTRUCTION, strategy=TruncationStrategy.NEVER_CUT),
            PriorityEntry(section=Section.OUTPUT_SPEC, strategy=TruncationStrategy.NEVER_CUT),
            PriorityEntry(section=Section.INPUT_DATA, strategy=TruncationStrategy.NEVER_CUT),
        ]),
    )

    # Tight budget to force truncation
    budget = TokenBudget(context_window=1_500, output_reserve=300, safety_margin=50)
    result = compile(prompt, budget)

    # ── Inspect what happened ──
    print(result.report())
    print()

    # Check which context blocks survived
    for section in result.sections:
        if section.section == Section.CONTEXT:
            print(f"Context: {section.items_included} kept, {section.items_dropped} dropped")
        if section.section == Section.HISTORY:
            print(f"History: {section.items_included} kept, {section.items_dropped} dropped")
        if section.section == Section.EXAMPLES:
            print(f"Examples: {section.items_included} kept, {section.items_dropped} dropped")

    print(f"\nHeadroom: {result.headroom} tokens")
    print(f"Warnings: {len(result.warnings)}")
    for w in result.warnings:
        print(f"  → {w}")

    return result.messages


if __name__ == "__main__":
    history = [
        {"role": "user", "content": "Hi, I have a question about my bill"},
        {"role": "assistant", "content": "Of course! I'd be happy to help with billing questions. What's going on?"},
        {"role": "user", "content": "I upgraded mid-month, will I be charged the full amount?"},
        {"role": "assistant", "content": "Great question — upgrades are prorated, so you'll only pay for the remaining days in your current billing cycle."},
    ]

    messages = build_chatbot_prompt(
        "Actually, I've decided I want to cancel instead. How do I do that?",
        history,
    )

    print("\n--- Final messages ---")
    for msg in messages:
        print(f"\n[{msg['role']}] {msg['content'][:150]}...")
