"""Gradio demo for the ai-steering package — offline/simulated mode.

No API key required and no live model calls are made. Each tab runs a real,
deterministic Python stand-in for the piece of logic that would otherwise
call Claude, so the demo is genuinely interactive without costing anything.
Install `ANTHROPIC_API_KEY` locally and run `examples/*.py` (or import
`ai_steering` directly) to see the real, Claude-backed behavior this Space
is illustrating.
"""

from __future__ import annotations

import re

import gradio as gr

from ai_steering.guardrails import GuardrailConfig, _keyword_prefilter
from ai_steering.routing import ROUTES

BANNER = (
    "> 🔌 **Simulated mode** — this Space makes no live model calls, so no API key is needed. "
    "Where noted, a small local heuristic stands in for the Claude call the real "
    "`ai_steering` function would make. The keyword prefilter in the Guardrails tab is the "
    "actual production code, since it's designed to never need a model call."
)

# ---------------------------------------------------------------------------
# Reasoning Enhancement — curated worked examples (the real chain_of_thought /
# tree_of_thoughts functions need a live Claude call; these are what they return).
# ---------------------------------------------------------------------------

_COT_EXAMPLES = {
    "A farmer has 17 sheep. All but 9 die. How many sheep are left?": {
        "thinking": (
            "'All but 9 die' means every sheep except 9 dies — so the number that survive "
            "is exactly the number excluded from dying, not 17 minus 9. The total of 17 is "
            "a distractor; the answer is the 9 explicitly spared."
        ),
        "answer": "Answer: 9",
    },
    "If a train travels 60 miles in 45 minutes, what is its speed in mph?": {
        "thinking": (
            "Speed = distance / time. Convert 45 minutes to hours: 45/60 = 0.75 hours. "
            "60 miles / 0.75 hours = 80 mph."
        ),
        "answer": "Answer: 80 mph",
    },
}

_TOT_EXAMPLE_BRANCHES = [
    {
        "approach": "Vertical (shard by tenant)",
        "reasoning": "Each tenant's data stays on one shard, keyed by tenant_id. Simple routing, no cross-shard joins for single-tenant queries.",
        "answer": "Shard by tenant_id with consistent hashing.",
        "score": 0.62,
        "critique": "Great isolation, but large tenants create hot shards and rebalancing is disruptive.",
    },
    {
        "approach": "Horizontal (shard by entity range)",
        "reasoning": "Split by a monotonic key range (e.g. user_id ranges). Enables even data growth distribution and simple range scans.",
        "answer": "Range-shard by user_id with a routing table for range-to-shard lookup.",
        "score": 0.78,
        "critique": "Better load distribution than tenant sharding, but range hotspots on sequential IDs need mitigation (e.g. hashed prefixes).",
    },
    {
        "approach": "Directory-based dynamic sharding",
        "reasoning": "A lookup service maps keys to shards explicitly, allowing shards to be split/merged/moved without a rehashing event.",
        "answer": "Directory service + hashed keys, with online shard splitting.",
        "score": 0.85,
        "critique": "Most operationally flexible and best long-term choice, at the cost of an extra directory-service hop and more moving parts to operate.",
    },
]


def run_reasoning(question: str, mode: str):
    if mode == "Chain of Thought":
        example = _COT_EXAMPLES.get(question) or next(iter(_COT_EXAMPLES.values()))
        return example["thinking"], example["answer"]

    branches_md = "\n\n".join(
        f"**{b['approach']}** (score {b['score']:.2f})\n{b['reasoning']}\n→ {b['answer']}\n\n*{b['critique']}*"
        for b in _TOT_EXAMPLE_BRANCHES
    )
    best = max(_TOT_EXAMPLE_BRANCHES, key=lambda b: b["score"])
    return branches_md, f"**Best answer:** {best['answer']}"


# ---------------------------------------------------------------------------
# Context Window Management — real buffer logic, with a word-count token
# estimate and a naive local condenser standing in for the Claude calls.
# ---------------------------------------------------------------------------

_context_messages: list[dict[str, str]] = []


def _estimate_tokens(messages: list[dict[str, str]]) -> int:
    return sum(len(m["content"].split()) for m in messages)


def _local_condense(messages: list[dict[str, str]]) -> str:
    words = " ".join(m["content"] for m in messages).split()
    preview = " ".join(words[:15])
    return f"{preview}... [{len(words)} words condensed into this summary]"


def context_reset(max_tokens: int, keep_recent: int):
    global _context_messages
    _context_messages = []
    return "Buffer reset.", 0


def context_add(turn: str, max_tokens: int, keep_recent: int):
    global _context_messages
    if not turn.strip():
        raise gr.Error("Enter a message first.")

    _context_messages.append({"role": "user", "content": turn})
    keep_recent = int(keep_recent)

    if _estimate_tokens(_context_messages) > int(max_tokens):
        recent = _context_messages[-keep_recent:] if keep_recent else []
        older = _context_messages[: len(_context_messages) - keep_recent] if keep_recent else _context_messages
        if older:
            summary = _local_condense(older)
            candidate = [
                {"role": "user", "content": f"[Earlier conversation summary — simulated]\n{summary}"},
                {"role": "assistant", "content": "Understood, I have the context from before."},
            ] + recent
            # If the condensed version still doesn't fit, drop the oldest kept-recent
            # messages one at a time (mirrors ai_steering.context.fit_to_budget).
            while len(candidate) > 2 and _estimate_tokens(candidate) > int(max_tokens):
                candidate.pop(2)
            _context_messages = candidate

    transcript = "\n\n".join(f"**{m['role']}**: {m['content']}" for m in _context_messages)
    return transcript, _estimate_tokens(_context_messages)


# ---------------------------------------------------------------------------
# Semantic Guardrails — the keyword prefilter is the real production code
# (it's designed to run without a model call); the semantic layer is a tiny
# word-overlap heuristic standing in for the Claude classifier call.
# ---------------------------------------------------------------------------

def run_guardrails(message: str, allowed_topics: str, blocked_keywords: str):
    if not message.strip():
        raise gr.Error("Enter a message first.")

    config = GuardrailConfig(
        allowed_topics=[t.strip() for t in allowed_topics.split(",") if t.strip()],
        blocked_keywords=[k.strip() for k in blocked_keywords.split(",") if k.strip()],
    )

    prefiltered = _keyword_prefilter(message, config)
    if prefiltered is not None:
        return f"🚫 **Blocked by keyword prefilter** (real code, no model call)\n\n{prefiltered.reason}"

    if config.allowed_topics:
        topic_words = {w.lower() for topic in config.allowed_topics for w in re.findall(r"\w+", topic)}
        message_words = {w.lower() for w in re.findall(r"\w+", message)}
        on_topic = bool(topic_words & message_words)
        if not on_topic:
            return (
                "🚫 **Blocked — off-topic** (simulated: word-overlap heuristic standing in for "
                f"the real semantic classifier)\n\nNo overlap with allowed topics: {', '.join(config.allowed_topics)}"
            )

    return (
        "✅ **Allowed** — keyword prefilter found nothing blocked, and the message is on-topic.\n\n"
        "*In live mode, the request would now go to Claude for a real answer.*"
    )


# ---------------------------------------------------------------------------
# Model Routing — a real local heuristic classifier standing in for the cheap
# Claude classification call; the routed answer itself is illustrative.
# ---------------------------------------------------------------------------

_COMPLEX_HINTS = ("design", "architecture", "fault-tolerant", "exactly-once", "distributed", "optimal", "trade-off", "tradeoff")
_MODERATE_HINTS = ("summarize", "compare", "explain", "why", "how")


def _local_classify(prompt: str) -> tuple[str, str]:
    lowered = prompt.lower()
    word_count = len(prompt.split())
    if any(h in lowered for h in _COMPLEX_HINTS) or word_count > 25:
        return "complex", "matched a complex-task keyword or the prompt is long"
    if any(h in lowered for h in _MODERATE_HINTS) or word_count > 10:
        return "moderate", "matched a multi-step-reasoning keyword or moderate length"
    return "simple", "short, direct factual request"


def run_routing(prompt: str):
    if not prompt.strip():
        raise gr.Error("Enter a prompt first.")
    complexity, reason = _local_classify(prompt)
    model = ROUTES[complexity]
    summary = f"Routed to **{model}** — classified as *{complexity}* (simulated heuristic: {reason})"
    return summary, "*In live mode, this is where the routed model's real answer would appear.*"


with gr.Blocks(title="AI Steering") as demo:
    gr.Markdown("# Core AI Steering Skills\n" + BANNER)

    with gr.Tab("Reasoning Enhancement"):
        gr.Markdown("Chain of Thought forces step-by-step reasoning; Tree of Thoughts explores several approaches and picks the best.")
        r_question = gr.Dropdown(
            list(_COT_EXAMPLES.keys()),
            value=next(iter(_COT_EXAMPLES.keys())),
            label="Question (Chain of Thought mode)",
        )
        r_mode = gr.Radio(["Chain of Thought", "Tree of Thoughts"], value="Chain of Thought", label="Mode")
        r_button = gr.Button("Show example")
        r_trace = gr.Markdown(label="Reasoning trace")
        r_answer = gr.Markdown(label="Answer")
        r_button.click(run_reasoning, inputs=[r_question, r_mode], outputs=[r_trace, r_answer])

    with gr.Tab("Context Window Management"):
        gr.Markdown("Add turns to a conversation with a small token budget and watch it condense older turns instead of forgetting them.")
        with gr.Row():
            c_budget = gr.Number(label="Max tokens (simulated, word-based)", value=25)
            c_keep = gr.Number(label="Keep recent messages", value=3)
        c_reset = gr.Button("Reset buffer")
        c_turn = gr.Textbox(label="Next message", value="My name is Priya and I'm building a rocket telemetry dashboard.")
        c_add = gr.Button("Add message")
        c_tokens = gr.Number(label="Current token count (simulated)", interactive=False)
        c_transcript = gr.Markdown(label="Buffered transcript")
        c_reset.click(context_reset, inputs=[c_budget, c_keep], outputs=[c_transcript, c_tokens])
        c_add.click(context_add, inputs=[c_turn, c_budget, c_keep], outputs=[c_transcript, c_tokens])

    with gr.Tab("Semantic Guardrails"):
        gr.Markdown("A keyword prefilter (real code) plus a semantic layer (simulated here) decide whether a message reaches the model.")
        g_topics = gr.Textbox(label="Allowed topics (comma-separated, blank = any topic)", value="billing subscription plan")
        g_keywords = gr.Textbox(label="Blocked keywords (comma-separated)", value="ignore previous instructions")
        g_message = gr.Textbox(label="User message", value="How do I upgrade my subscription plan?")
        g_button = gr.Button("Check")
        g_output = gr.Markdown()
        g_button.click(run_guardrails, inputs=[g_message, g_topics, g_keywords], outputs=g_output)

    with gr.Tab("Model Routing"):
        gr.Markdown("A cheap classifier (simulated here as a local heuristic) decides which Claude tier a request is routed to.")
        m_prompt = gr.Textbox(label="Prompt", value="What's the capital of Japan?")
        m_button = gr.Button("Route")
        m_summary = gr.Markdown()
        m_answer = gr.Markdown()
        m_button.click(run_routing, inputs=[m_prompt], outputs=[m_summary, m_answer])

if __name__ == "__main__":
    demo.launch()