"""
app.py  –  Enterprise SQL Agent (Gradio + smolagents + MCP)

SECRETS / ENV VARS
------------------
OPENAI_API_KEY   ← use OpenAI (default model gpt-4o, override with OPENAI_MODEL)
GOOGLE_API_KEY   ← use Gemini-Pro (override model with GOOGLE_MODEL)
HF_MODEL_ID      ← repo that exposes Chat-Completion (fallback if no keys)
HF_API_TOKEN     ← token if that repo is gated

FILE LAYOUT
-----------
app.py
mcp_server.py      # your FastMCP SQL tool server
requirements.txt   # see bottom of this file
"""

import os, pathlib, gradio as gr
from mcp import StdioServerParameters
from smolagents import MCPClient, CodeAgent
from smolagents.models import LiteLLMModel, InferenceClientModel

# ─────────── 1. Choose base LLM ──────────────────────────────────────────
OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")          # change if not whitelisted

GEMINI_KEY   = os.getenv("GOOGLE_API_KEY")
GEM_MODEL    = os.getenv("GOOGLE_MODEL", "gemini-pro")

HF_MODEL_ID  = os.getenv("HF_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
HF_TOKEN     = os.getenv("HF_API_TOKEN")  # only for gated repos

if OPENAI_KEY:
    BASE_MODEL = LiteLLMModel(model_id=f"openai/{OPENAI_MODEL}", api_key=OPENAI_KEY)
    ACTIVE     = f"OpenAI · {OPENAI_MODEL}"
elif GEMINI_KEY:
    BASE_MODEL = LiteLLMModel(model_id=f"google/{GEM_MODEL}", api_key=GEMINI_KEY)
    ACTIVE     = f"Gemini · {GEM_MODEL}"
else:
    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, hf_api_token=HF_TOKEN)
    ACTIVE     = f"Hugging Face · {HF_MODEL_ID}"

# ─────────── 2. Path to MCP tool server ──────────────────────────────────
SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")

# ─────────── 3. Gradio callback ──────────────────────────────────────────
def respond(msg: str, history: list):
    """Run prompt → CodeAgent → MCP tools → safe string reply."""
    params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])

    with MCPClient(params) as tools:
        agent  = CodeAgent(tools=tools, model=BASE_MODEL)
        raw    = agent.run(msg)

    # Ensure reply is always string for Chatbot
    if not isinstance(raw, str):
        import json, pprint
        try:
            raw = json.dumps(raw, indent=2, ensure_ascii=False)
        except (TypeError, ValueError):
            raw = pprint.pformat(raw)
    reply = raw

    history += [
        {"role": "user",      "content": msg},
        {"role": "assistant", "content": reply},
    ]
    return history, history

# ─────────── 4. Build the UI ─────────────────────────────────────────────
with gr.Blocks(title="Enterprise SQL Agent") as demo:
    state = gr.State([])
    gr.Markdown("## 🏢 Enterprise SQL Agent — ask natural-language questions about your data")

    chat = gr.Chatbot(type="messages", label="Conversation")
    box  = gr.Textbox(
        placeholder="e.g. Who are my Northeast customers with no orders in 6 months?",
        show_label=False,
    )
    box.submit(respond, [box, state], [chat, state])

    with gr.Accordion("Example prompts", open=False):
        gr.Markdown(
            "* Who are my **Northeast** customers with no orders in 6 months?\n"
            "* List customers sorted by **LastOrderDate**.\n"
            "* Draft re-engagement emails for inactive accounts."
        )

    gr.Markdown(f"_Powered by MCP + smolagents + Gradio • Active model → **{ACTIVE}**_")

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