import gradio as gr
from typing import List, Dict, Set, Callable, Any
import json
import uuid
from pathlib import Path

# Azure SDK auth (Entra ID)
from azure.identity import DefaultAzureCredential

# Azure AI Agents SDK
from azure.ai.agents import AgentsClient
from azure.ai.agents.models import (
    FunctionTool,
    ToolSet,
    ListSortOrder,
    MessageRole,
)

# ---------- Inline custom function(s) ----------

def submit_support_ticket(email_address: str, description: str) -> str:
    """
    Generates a ticket number and saves a support ticket as a text file
    in the current app directory. Returns a small JSON message string.
    """
    script_dir = Path(__file__).parent
    ticket_number = str(uuid.uuid4()).replace("-", "")[:6]
    file_name = f"ticket-{ticket_number}.txt"
    file_path = script_dir / file_name

    text = (
        f"Support ticket: {ticket_number}\n"
        f"Submitted by: {email_address}\n"
        f"Description:\n{description}\n"
    )
    file_path.write_text(text, encoding="utf-8")

    return json.dumps({
        "message": f"Support ticket {ticket_number} submitted. The ticket file is saved as {file_name}"
    })

# Register callable tools
user_functions: Set[Callable[..., Any]] = { submit_support_ticket }

# ---------- Core Agent Helpers ----------

def init_agent(project_endpoint: str, model_deployment: str) -> dict:
    """
    Initialize an Azure AI Agent (Entra ID auth via DefaultAzureCredential).
    Returns a session dict containing client, agent_id, thread_id, etc.
    """
    if not project_endpoint or not model_deployment:
        raise ValueError("Please provide a Project Endpoint and Model Deployment name (e.g., gpt-4o).")

    # Entra ID token credential (requires `az login` or other supported auth in your environment)
    credential = DefaultAzureCredential(
        exclude_environment_credential=False,
        exclude_managed_identity_credential=False,
        exclude_shared_token_cache_credential=False,
        exclude_visual_studio_code_credential=False,
        exclude_powershell_credential=False,
        exclude_cli_credential=False,  # allows Azure CLI token if you've run `az login`
        exclude_interactive_browser_credential=True,  # set True for server environments
    )

    client = AgentsClient(
        endpoint=project_endpoint.strip(),
        credential=credential,
    )

    # Build toolset and enable auto function calls
    functions_tool = FunctionTool(user_functions)
    toolset = ToolSet()
    toolset.add(functions_tool)
    client.enable_auto_function_calls(toolset)

    # Create the agent
    agent = client.create_agent(
        model=model_deployment.strip(),
        name="support-agent",
        instructions=(
            "You are a technical support agent. "
            "Ask for the user's email address and a description of the issue. "
            "Then submit a support ticket using the provided function. "
            "If a file is saved, tell the user the file name."
        ),
        toolset=toolset,
    )

    # Create a thread for the conversation
    thread = client.threads.create()

    # Session we keep in Gradio state
    return {
        "endpoint": project_endpoint.strip(),
        "model": model_deployment.strip(),
        "client": client,
        "agent_id": agent.id,
        "thread_id": thread.id,
    }


def send_to_agent(user_msg: str, session: dict):
    """
    Send message to the agent thread and return:
    - agent_reply (str)
    - history_str (str) chronological log
    """
    if not session or "client" not in session:
        raise ValueError("Agent is not initialized. Click 'Connect & Prepare' first.")

    client: AgentsClient = session["client"]
    agent_id = session["agent_id"]
    thread_id = session["thread_id"]

    # Add user message
    client.messages.create(
        thread_id=thread_id,
        role="user",
        content=user_msg,
    )

    # Run and wait (auto function-calls enabled)
    run = client.runs.create_and_process(thread_id=thread_id, agent_id=agent_id)
    if getattr(run, "status", None) == "failed":
        last_error = getattr(run, "last_error", "Unknown error")
        return f"Run failed: {last_error}", ""

    # Last agent message
    last_msg = client.messages.get_last_message_text_by_role(
        thread_id=thread_id,
        role=MessageRole.AGENT,
    )
    agent_reply = last_msg.text.value if last_msg else "(No reply text found.)"

    # Build readable history (chronological)
    history_lines = []
    messages = client.messages.list(thread_id=thread_id, order=ListSortOrder.ASCENDING)
    for m in messages:
        if m.text_messages:
            last_text = m.text_messages[-1].text.value
            history_lines.append(f"{m.role}: {last_text}")
    history_str = "\n\n".join(history_lines)

    return agent_reply, history_str


def teardown(session: dict) -> str:
    """
    Delete the agent to reduce costs. (Threads are retained by the service.)
    """
    if not session:
        return "Nothing to clean up."

    notes = []
    try:
        client: AgentsClient = session.get("client")
        agent_id = session.get("agent_id")
        if client and agent_id:
            client.delete_agent(agent_id)
            notes.append("Deleted agent.")
    except Exception as e:
        notes.append(f"Cleanup warning: {e}")

    return " ".join(notes) if notes else "Cleanup complete."


# ---------- Gradio UI ----------

with gr.Blocks(title="Azure AI Support Agent (Functions) — Gradio") as demo:
    gr.Markdown(
        "## Azure AI Support Agent (Custom Function Tool)\n"
        "1) **Run `az login`** in the same environment.  \n"
        "2) Paste your **Project Endpoint** and **Model Deployment** (e.g., `gpt-4o`).  \n"
        "Then chat — the agent can auto-call a function to submit a support ticket."
    )

    with gr.Row():
        endpoint = gr.Textbox(label="Project Endpoint", placeholder="https://<your-project-endpoint>")
        model = gr.Textbox(label="Model Deployment Name", value="gpt-4o")

    session_state = gr.State(value=None)

    connect_btn = gr.Button("🔌 Connect & Prepare Agent", variant="primary")
    connect_status = gr.Markdown("")

    with gr.Row():
        chatbot = gr.Chatbot(
            label="Conversation",
            height=420,
            type="messages",  # openai-style dicts: {"role": "...", "content": "..."}
        )

    user_input = gr.Textbox(label="Your message", placeholder="e.g., I have a technical problem")
    with gr.Row():
        send_btn = gr.Button("Send ▶")
        cleanup_btn = gr.Button("Delete Agent & Cleanup 🧹")

    history = gr.Textbox(label="Conversation Log (chronological)", lines=12)

    # ----- Callbacks -----

    def on_connect(ep, mdl):
        try:
            sess = init_agent(ep, mdl)
            return sess, "✅ Connected. Support agent and thread are ready. (Auth via DefaultAzureCredential)"
        except Exception as e:
            return None, f"❌ Connection error: {e}"

    connect_btn.click(
        fn=on_connect,
        inputs=[endpoint, model],
        outputs=[session_state, connect_status],
    )

    def on_send(msg: str, session: dict, chat_msgs: List[Dict[str, str]]):
        if not msg:
            return gr.update(), gr.update(), gr.update(value="Please enter a message.")
        try:
            agent_reply, log = send_to_agent(msg, session)
            chat_msgs = (chat_msgs or []) + [
                {"role": "user", "content": msg},
                {"role": "assistant", "content": agent_reply},
            ]
            return chat_msgs, "", gr.update(value=log)  # clear input, update log
        except Exception as e:
            return chat_msgs, msg, gr.update(value=f"❌ Error: {e}")

    send_btn.click(
        fn=on_send,
        inputs=[user_input, session_state, chatbot],
        outputs=[chatbot, user_input, history],
    )

    def on_cleanup(session):
        try:
            msg = teardown(session)
            return None, f"🧹 {msg}"
        except Exception as e:
            return session, f"⚠️ Cleanup error: {e}"

    cleanup_btn.click(
        fn=on_cleanup,
        inputs=[session_state],
        outputs=[session_state, connect_status],
    )

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