"""
Run locally:
  python -m pip install --upgrade pip
  pip install gradio scikit-learn numpy pandas
  python RAG_vs_MCP_Demo_App.py


"""
from __future__ import annotations

import re
import json
from dataclasses import dataclass, asdict
from typing import Any, Callable, Dict, List, Tuple, Optional

import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr

# ------------------------------
# 1) Tiny knowledge base (handbook) for RAG
# ------------------------------
DOCS = [
    {
        "id": "policy_returns",
        "title": "Returns & Refunds Policy",
        "text": (
            "Customers can return most new, unopened items within 30 days of delivery for a full refund. "
            "Items must be in their original condition with receipt. Refunds are processed to the original payment method. "
            "Defective or damaged items are eligible for free return shipping."
        ),
    },
    {
        "id": "policy_shipping",
        "title": "Shipping Policy",
        "text": (
            "Standard shipping typically takes 3 to 5 business days. Expedited shipping options are available at checkout. "
            "International orders may take 7 to 14 business days depending on destination and customs."
        ),
    },
    {
        "id": "policy_membership",
        "title": "Membership Benefits",
        "text": (
            "Members earn 2 points per dollar spent, get early access to new releases, and receive a monthly newsletter with curated picks. "
            "Points can be redeemed for discounts at checkout."
        ),
    },
    {
        "id": "policy_giftcards",
        "title": "Gift Cards",
        "text": (
            "Gift cards are available in denominations from $10 to $200 and are redeemable online or in-store. "
            "They do not expire and cannot be redeemed for cash except where required by law."
        ),
    },
]

# Fit a very small TF‑IDF retriever at startup
VECTORIZER = TfidfVectorizer(stop_words="english")
KB_TEXTS = [d["text"] for d in DOCS]
KB_MATRIX = VECTORIZER.fit_transform(KB_TEXTS)


def rag_retrieve(query: str, k: int = 3) -> List[Dict[str, Any]]:
    """Return top-k documents as {id,title,text,score}."""
    if not query.strip():
        return []
    q_vec = VECTORIZER.transform([query])
    sims = cosine_similarity(q_vec, KB_MATRIX)[0]
    idxs = np.argsort(-sims)[:k]
    results = []
    for i in idxs:
        results.append({
            "id": DOCS[i]["id"],
            "title": DOCS[i]["title"],
            "text": DOCS[i]["text"],
            "score": float(sims[i]),
        })
    return results


def rag_answer(query: str, k: int = 3) -> Tuple[str, List[Dict[str, Any]]]:
    """Simple, template-y answer based on top-k docs."""
    hits = rag_retrieve(query, k=k)
    if not hits:
        return ("I couldn't find anything relevant in the handbook.", [])

    # Compose a short grounded answer using snippets
    bullets = []
    for h in hits:
        # Take the first sentence as a snippet
        first_sentence = h["text"].split(".")[0].strip()
        if first_sentence:
            bullets.append(f"- **{h['title']}**: {first_sentence}.")
    answer = (
        "**Handbook says:**\n" + "\n".join(bullets) +
        "\n\n(Answer generated from retrieved policy snippets; no LLM used.)"
    )
    return answer, hits

# ------------------------------
# 2) MCP-style tool registry + client executor
# ------------------------------
@dataclass
class ToolParam:
    name: str
    type: str  # e.g., "string", "number", "integer"
    description: str
    required: bool = True


@dataclass
class ToolSpec:
    name: str
    description: str
    params: List[ToolParam]


@dataclass
class ToolCall:
    tool_name: str
    args: Dict[str, Any]


@dataclass
class ToolResult:
    tool_name: str
    args: Dict[str, Any]
    result: Any
    ok: bool
    error: Optional[str] = None


# In-memory "live" inventory
INVENTORY: Dict[str, Dict[str, Any]] = {
    "Dune": {"stock": 7, "price": 19.99},
    "Clean Code": {"stock": 2, "price": 25.99},
    "The Pragmatic Programmer": {"stock": 5, "price": 31.50},
    "Deep Learning": {"stock": 1, "price": 64.00},
}


# Define actual tool functions
def tool_get_inventory(title: str) -> Dict[str, Any]:
    rec = INVENTORY.get(title)
    if not rec:
        return {"title": title, "found": False, "message": f"'{title}' not in inventory."}
    return {"title": title, "found": True, **rec}


def tool_set_price(title: str, new_price: float) -> Dict[str, Any]:
    rec = INVENTORY.get(title)
    if not rec:
        return {"title": title, "updated": False, "message": f"'{title}' not in inventory."}
    rec["price"] = float(new_price)
    return {"title": title, "updated": True, **rec}


def tool_place_order(title: str, quantity: int) -> Dict[str, Any]:
    rec = INVENTORY.get(title)
    if not rec:
        return {"title": title, "ordered": False, "message": f"'{title}' not in inventory."}
    if quantity <= 0:
        return {"title": title, "ordered": False, "message": "Quantity must be positive."}
    rec["stock"] += int(quantity)
    return {"title": title, "ordered": True, "added": int(quantity), **rec}


# Registry of specs (like MCP manifests)
TOOL_SPECS: Dict[str, ToolSpec] = {
    "get_inventory": ToolSpec(
        name="get_inventory",
        description="Get stock and price for a given book title.",
        params=[
            ToolParam("title", "string", "Exact book title"),
        ],
    ),
    "set_price": ToolSpec(
        name="set_price",
        description="Update the price for a book title.",
        params=[
            ToolParam("title", "string", "Exact book title"),
            ToolParam("new_price", "number", "New price in dollars"),
        ],
    ),
    "place_order": ToolSpec(
        name="place_order",
        description="Increase stock by ordering more copies.",
        params=[
            ToolParam("title", "string", "Exact book title"),
            ToolParam("quantity", "integer", "How many copies to add"),
        ],
    ),
}

# Mapping tool names to callables
TOOL_IMPLS: Dict[str, Callable[..., Any]] = {
    "get_inventory": tool_get_inventory,
    "set_price": tool_set_price,
    "place_order": tool_place_order,
}


def validate_and_call(call: ToolCall) -> ToolResult:
    spec = TOOL_SPECS.get(call.tool_name)
    if not spec:
        return ToolResult(tool_name=call.tool_name, args=call.args, result=None, ok=False, error="Unknown tool")

    # minimal validation
    for p in spec.params:
        if p.required and p.name not in call.args:
            return ToolResult(tool_name=call.tool_name, args=call.args, result=None, ok=False, error=f"Missing param: {p.name}")

    try:
        fn = TOOL_IMPLS[call.tool_name]
        result = fn(**call.args)
        return ToolResult(tool_name=call.tool_name, args=call.args, result=result, ok=True)
    except Exception as e:
        return ToolResult(tool_name=call.tool_name, args=call.args, result=None, ok=False, error=str(e))

# ------------------------------
# 3) Simple planner/router: choose RAG vs Tools (MCP-style) vs Both
# ------------------------------
TOOL_KEYWORDS = {
    "get_inventory": ["in stock", "stock", "available", "availability", "have", "inventory"],
    "set_price": ["change price", "set price", "update price", "price to", "discount", "mark down"],
    "place_order": ["order", "restock", "add", "increase stock"],
}

BOOK_TITLE_PATTERN = r"'([^']+)'|\"([^\"]+)\""  # capture 'Title' or "Title"


def extract_titles(text: str) -> List[str]:
    titles = []
    for m in re.finditer(BOOK_TITLE_PATTERN, text):
        titles.append(m.group(1) or m.group(2))
    return titles


def decide_tools(query: str) -> Optional[ToolCall]:
    q = query.lower()
    titles = extract_titles(query)

    # get_inventory
    if any(kw in q for kw in TOOL_KEYWORDS["get_inventory"]):
        if titles:
            return ToolCall(tool_name="get_inventory", args={"title": titles[0]})

    # set_price  (look for a number)
    if any(kw in q for kw in TOOL_KEYWORDS["set_price"]):
        price_match = re.search(r"(\d+\.?\d*)", q)
        if titles and price_match:
            return ToolCall(tool_name="set_price", args={"title": titles[0], "new_price": float(price_match.group(1))})

    # place_order  (look for an integer quantity)
    if any(kw in q for kw in TOOL_KEYWORDS["place_order"]):
        qty_match = re.search(r"(\d+)", q)
        if titles and qty_match:
            return ToolCall(tool_name="place_order", args={"title": titles[0], "quantity": int(qty_match.group(1))})

    return None


def route_query(query: str, mode: str = "Auto") -> str:
    if mode == "RAG only":
        return "rag"
    if mode == "Tools only":
        return "tools"

    # Auto: detect whether we need tools, rag, or both
    # If a single sentence includes both a policy question + inventory check, we'll call it "both".
    needs_tool = decide_tools(query) is not None
    needs_rag = any(ch in query.lower() for ch in ["policy", "return", "refund", "shipping", "membership", "gift card", "gift cards", "benefits"])

    if needs_tool and needs_rag:
        return "both"
    if needs_tool:
        return "tools"
    return "rag"

# ------------------------------
# 4) Orchestrator: build a human-friendly answer + trace
# ------------------------------

def handle_query(query: str, mode: str = "Auto", show_trace: bool = True) -> Tuple[str, str, pd.DataFrame]:
    route = route_query(query, mode=mode)

    tool_trace: List[Dict[str, Any]] = []
    rag_hits: List[Dict[str, Any]] = []
    parts: List[str] = []

    if route in ("rag", "both"):
        rag_ans, rag_hits = rag_answer(query)
        parts.append(rag_ans)

    if route in ("tools", "both"):
        call = decide_tools(query)
        if call:
            res = validate_and_call(call)
            tool_trace.append(asdict(call))
            tool_trace[-1]["result"] = res.result
            tool_trace[-1]["ok"] = res.ok
            if res.error:
                tool_trace[-1]["error"] = res.error

            # Compose a user-friendly tool result string
            if res.ok and isinstance(res.result, dict):
                if call.tool_name == "get_inventory":
                    if res.result.get("found"):
                        parts.append(
                            f"**Inventory:** '{res.result['title']}' -- stock: {res.result['stock']}, price: ${res.result['price']:.2f}"
                        )
                    else:
                        parts.append(f"**Inventory:** {res.result.get('message','Not found')}" )
                elif call.tool_name == "set_price":
                    if res.result.get("updated"):
                        parts.append(
                            f"**Price updated:** '{res.result['title']}' is now ${res.result['price']:.2f}"
                        )
                    else:
                        parts.append(f"**Set price failed:** {res.result.get('message','Error')}" )
                elif call.tool_name == "place_order":
                    if res.result.get("ordered"):
                        parts.append(
                            f"**Order placed:** Added {res.result['added']} copies of '{res.result['title']}'. New stock: {res.result['stock']}"
                        )
                    else:
                        parts.append(f"**Order failed:** {res.result.get('message','Error')}" )
            else:
                parts.append("Tool call failed.")
        else:
            parts.append("No suitable tool call inferred from your request.")

    # Prepare trace artifacts
    trace = {
        "route": route,
        "tool_calls": tool_trace,
        "retrieved_docs": rag_hits,
    }

    # DataFrame for retrieved docs (for a quick visual)
    df = pd.DataFrame([
        {
            "id": h["id"],
            "title": h["title"],
            "score": round(h["score"], 3),
            "snippet": h["text"][:140] + ("..." if len(h["text"])>140 else ""),
        }
        for h in rag_hits
    ])

    answer_md = "\n\n".join(parts) if parts else "(No answer composed.)"
    trace_json = json.dumps(trace, indent=2)

    return answer_md, trace_json, df

# ------------------------------
# 5) Gradio UI
# ------------------------------
with gr.Blocks(title="RAG vs MCP Demo: Bookstore Ops Assistant") as demo:
    gr.Markdown(
        "# RAG vs MCP Demo: Bookstore Ops Assistant\n"
        "Use this sandbox to feel the difference between RAG (lookup from a handbook) and MCP-style tools (act on live data).\n\n"
        "**Tips**: Put book titles in quotes, e.g., 'Dune' or \"Clean Code\"."
    )

    with gr.Row():
        query = gr.Textbox(label="Your request", placeholder="e.g., Do we have 'Dune' in stock? Or: What is our returns policy?", lines=2)
    with gr.Row():
        mode = gr.Dropdown(["Auto", "RAG only", "Tools only"], value="Auto", label="Routing mode")
        show_trace = gr.Checkbox(True, label="Show trace")
        submit = gr.Button("Run", variant="primary")

    answer = gr.Markdown(label="Answer")
    trace = gr.JSON(label="Trace (route, tool calls, retrieved docs)")
    table = gr.Dataframe(headers=["id", "title", "score", "snippet"], label="Retrieved docs (RAG)")

    def _run(q, m, t):
        ans, tr, df = handle_query(q or "", mode=m or "Auto", show_trace=bool(t))
        return ans, json.loads(tr), df

    submit.click(_run, inputs=[query, mode, show_trace], outputs=[answer, trace, table])

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