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

ENV VARS SUPPORTED
------------------
OPENAI_API_KEY   → use any OpenAI chat model (default  = gpt-4o)
OPENAI_MODEL     → override the OpenAI model ID   (e.g. gpt-4-turbo)
GOOGLE_API_KEY   → use Gemini-Pro via Google AI
GOOGLE_MODEL     → override Gemini model ID      (e.g. gemini-1.5-pro)

HF_MODEL_ID      → fallback model repo (must expose Chat-Completion)
HF_API_TOKEN     → token if that repo is gated

If no provider keys are set, the code falls back to a free HF model
`microsoft/Phi-3-mini-4k-instruct`.

MCP SERVER
----------
`mcp_server.py` must live in the same folder and expose your SQL tools.
"""

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

# --------------------------------------------------------------------- #
# 1.  Resolve the base LLM model, preferring OpenAI → Gemini → HF
# --------------------------------------------------------------------- #
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") or None            # only needed 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 your MCP server (must be alongside this file)
# --------------------------------------------------------------------- #
SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")

# --------------------------------------------------------------------- #
# 3.  Gradio callback – run prompt → CodeAgent → SQL tools
# --------------------------------------------------------------------- #
def respond(msg: str, chat_history: list):
    params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])

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

    chat_history += [
        {"role": "user",      "content": msg},
        {"role": "assistant", "content": reply},
    ]
    return chat_history, chat_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="Chat")
    box    = gr.Textbox(
        show_label=False,
        placeholder="e.g. Who are my Northeast customers with no orders in 6 months?",
    )
    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 tools + smolagents · Active model → **{ACTIVE}**_")

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