"""
Agent Ken — Data-Informed PM + AI Copilot
Powered by Azure AI Foundry + Microsoft Fabric
HuggingFace Space (Gradio) — with file upload support
"""

import os
import gradio as gr
from azure.ai.projects import AIProjectClient
from azure.identity import ClientSecretCredential
from azure.ai.agents.models import ListSortOrder, FilePurpose, MessageAttachment, CodeInterpreterTool

# Azure config
PROJECT_ENDPOINT = os.environ["PROJECT_ENDPOINT"]
AGENT_ID = os.environ["AGENT_ID"]
AZURE_TENANT_ID = os.environ["AZURE_TENANT_ID"]
AZURE_CLIENT_ID = os.environ["AZURE_CLIENT_ID"]
AZURE_CLIENT_SECRET = os.environ["AZURE_CLIENT_SECRET"]

# Azure client
credential = ClientSecretCredential(
    tenant_id=AZURE_TENANT_ID,
    client_id=AZURE_CLIENT_ID,
    client_secret=AZURE_CLIENT_SECRET,
)
project = AIProjectClient(credential=credential, endpoint=PROJECT_ENDPOINT)

_agent = None

def get_agent():
    global _agent
    if _agent is None:
        _agent = project.agents.get_agent(AGENT_ID)
    return _agent


# Supported file types
SUPPORTED_EXTENSIONS = {
    ".csv", ".xlsx", ".xls", ".json",
    ".txt", ".md", ".pdf",
    ".doc", ".docx", ".pptx",
    ".png", ".jpg", ".jpeg", ".gif",
    ".py", ".html", ".css", ".js",
    ".xml", ".zip", ".tar",
}

def is_supported_file(filename: str) -> bool:
    ext = os.path.splitext(filename)[1].lower()
    return ext in SUPPORTED_EXTENSIONS

# Chat logic (generator for typing indicator)
def respond(user_message: str, uploaded_files, history: list, thread_state: dict | None):
    # Normalize inputs
    user_message = (user_message or "").strip()
    if not uploaded_files:
        uploaded_files = []

    # Nothing to send
    if not user_message and not uploaded_files:
        yield history, thread_state
        return

    # Create thread if needed
    if thread_state is None or "thread_id" not in thread_state:
        thread = project.agents.threads.create()
        thread_state = {"thread_id": thread.id}

    thread_id = thread_state["thread_id"]

    # Show typing indicator immediately
    display_msg = user_message
    if uploaded_files:
        file_names_display = [os.path.basename(str(f)) for f in uploaded_files]
        if user_message:
            display_msg = f"{user_message}\n\n📎 {', '.join(file_names_display)}"
        else:
            display_msg = f"📎 {', '.join(file_names_display)}"

    history.append({"role": "user", "content": display_msg})
    history.append({"role": "assistant", "content": "🧠 Thinking..."})
    yield history, thread_state

    try:
        agent = get_agent()

        # Handle file uploads
        attachments = []
        file_names = []

        for file_path in uploaded_files:
            path_str = str(file_path)
            filename = os.path.basename(path_str)

            if not os.path.exists(path_str):
                history[-1] = {
                    "role": "assistant",
                    "content": f"⚠️ File `{filename}` could not be read. Please re-upload and try again."
                }
                yield history, thread_state
                return

            if not is_supported_file(filename):
                history[-1] = {
                    "role": "assistant",
                    "content": (
                        f"⚠️ File `{filename}` is not supported.\n\n"
                        f"**Supported formats:** CSV, Excel, PDF, Word, TXT, PNG, JPG, JSON, PowerPoint, Python, HTML, XML."
                    )
                }
                yield history, thread_state
                return

            try:
                uploaded = project.agents.files.upload_and_poll(
                    file_path=path_str,
                    purpose=FilePurpose.AGENTS,
                )
                attachments.append(
                    MessageAttachment(
                        file_id=uploaded.id,
                        tools=[CodeInterpreterTool().definitions[0]],
                    )
                )
                file_names.append(filename)
            except Exception as upload_err:
                history[-1] = {
                    "role": "assistant",
                    "content": f"⚠️ Failed to upload `{filename}`: {str(upload_err)}\n\nPlease try again."
                }
                yield history, thread_state
                return

        # Build message content
        if file_names and user_message:
            content = (
                f"{user_message}\n\n"
                f"📎 Uploaded files: {', '.join(file_names)}\n\n"
                f"Instructions: Read and analyze the uploaded file(s) immediately using code interpreter. "
                f"Show the data structure, key findings, and actionable insights. "
                f"Do NOT ask the user to re-upload or confirm — the file is already attached and accessible."
            )
        elif file_names:
            content = (
                f"📎 Uploaded files: {', '.join(file_names)}\n\n"
                f"Instructions: Read and analyze the uploaded file(s) immediately using code interpreter. "
                f"Show a summary of the data (columns, rows, data types), then provide key findings, "
                f"trends, and actionable insights. Start analyzing right away."
            )
        else:
            content = user_message

        # Send message to Azure
        msg_kwargs = {
            "thread_id": thread_id,
            "role": "user",
            "content": content,
        }
        if attachments:
            msg_kwargs["attachments"] = attachments

        project.agents.messages.create(**msg_kwargs)

        # Run agent
        run = project.agents.runs.create_and_process(
            thread_id=thread_id,
            agent_id=agent.id,
        )

        if run.status == "failed":
            assistant_reply = (
                f"⚠️ Agent run failed: {run.last_error}\n\n"
                "Please try again or start a new conversation."
            )
        else:
            messages = project.agents.messages.list(
                thread_id=thread_id, order=ListSortOrder.DESCENDING
            )
            assistant_reply = "🤔 No response received. Please try again."
            for msg in messages:
                if msg.role == "assistant" and msg.text_messages:
                    assistant_reply = msg.text_messages[-1].text.value
                    break

    except Exception as e:
        assistant_reply = (
            f"❌ Connection error: {str(e)}\n\n"
            "Please try again. If this persists, the Azure endpoint may be temporarily unavailable."
        )

    # Replace typing indicator with actual response
    history[-1] = {"role": "assistant", "content": assistant_reply}
    yield history, thread_state


def new_conversation():
    return [], None


# CSS
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

* {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

.gradio-container {
    max-width: 880px !important;
    margin: 0 auto !important;
}

footer { display: none !important; }

/* Header */
.header-main {
    text-align: center;
    padding: 28px 16px 6px 16px;
}
.header-main h1 {
    font-size: 1.7rem;
    font-weight: 700;
    color: #1A202C;
    margin: 0 0 4px 0;
}
.header-main .tagline {
    font-size: 0.88rem;
    color: #718096;
    margin: 0;
}


/* Architecture badge */
.arch-badge {
    display: flex;
    justify-content: center;
    gap: 8px;
    flex-wrap: wrap;
    margin: 14px auto 6px auto;
    max-width: 640px;
}
.arch-badge .chip {
    display: inline-flex;
    align-items: center;
    gap: 5px;
    padding: 5px 13px;
    border-radius: 8px;
    font-size: 0.72rem;
    font-weight: 600;
    letter-spacing: 0.02em;
}
.chip-foundry {
    background: #EBF4FF;
    color: #2B6CB0;
    border: 1px solid #BEE3F8;
}
.chip-fabric {
    background: #F0FFF4;
    color: #276749;
    border: 1px solid #C6F6D5;
}
.chip-ml {
    background: #FAF5FF;
    color: #6B46C1;
    border: 1px solid #E9D8FD;
}
.chip-model {
    background: #FFFAF0;
    color: #C05621;
    border: 1px solid #FEEBC8;
}

/* Welcome card */
.welcome-card {
    background: linear-gradient(135deg, #EBF4FF 0%, #F7FAFC 40%, #F0FFF4 100%);
    border: 1px solid #D6E4F0;
    border-radius: 16px;
    padding: 24px 28px 20px 28px;
    margin: 10px 0 14px 0;
    text-align: center;
}
.welcome-card .wave {
    font-size: 2rem;
    display: block;
    margin-bottom: 6px;
}
.welcome-card h2 {
    font-size: 1.1rem;
    font-weight: 600;
    color: #2D3748;
    margin: 0 0 6px 0;
}
.welcome-card p {
    font-size: 0.84rem;
    color: #4A5568;
    margin: 0 0 14px 0;
    line-height: 1.55;
}
.welcome-tags {
    display: flex;
    flex-wrap: wrap;
    gap: 6px;
    justify-content: center;
}
.welcome-tags span {
    background: white;
    border: 1px solid #E2E8F0;
    border-radius: 20px;
    padding: 5px 13px;
    font-size: 0.74rem;
    color: #4A5568;
    font-weight: 500;
}

/* Data section */
.data-section {
    background: #F7FAFC;
    border: 1px solid #E2E8F0;
    border-radius: 12px;
    padding: 14px 20px;
    margin: 0 0 14px 0;
}
.data-section h3 {
    font-size: 0.82rem;
    font-weight: 600;
    color: #2D3748;
    margin: 0 0 8px 0;
}
.data-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
    gap: 8px;
}
.data-card {
    background: white;
    border: 1px solid #E2E8F0;
    border-radius: 10px;
    padding: 10px 12px;
    text-align: center;
}
.data-card .number {
    font-size: 1.15rem;
    font-weight: 700;
    color: #2B6CB0;
    display: block;
}
.data-card .label {
    font-size: 0.68rem;
    color: #718096;
    margin-top: 2px;
    display: block;
}

/* Upload hint */
.upload-hint {
    background: #FFFFF0;
    border: 1px solid #FEFCBF;
    border-radius: 10px;
    padding: 10px 16px;
    margin: 0 0 10px 0;
    font-size: 0.76rem;
    color: #744210;
    text-align: center;
}

/* Footer */
.footer-info {
    text-align: center;
    padding: 10px 0;
    font-size: 0.7rem;
    color: #A0AEC0;
    line-height: 1.6;
}
"""

# Examples
EXAMPLES = [
    ["What's our D7 retention rate?"],
    ["Which acquisition channel has the highest LTV?"],
    ["How did our checkout experiment perform?"],
    ["Which user segments are at highest churn risk?"],
    ["Help me write a PRD for a referral program feature"],
    ["Score these features using RICE: push notifications, dark mode, onboarding revamp"],
    ["Design an A/B test for our new premium pricing"],
    ["Were there any anomalies in our metrics recently?"],
]

# UI
with gr.Blocks(title="Agent Ken — Data-Informed PM Copilot") as demo:

    # Header
    gr.HTML("""
        <div class="header-main">
            <h1>🤖 Agent Ken</h1>
            <p class="tagline">Your AI Companion for Product Management & Data</p>
        </div>
    """)

    # Architecture chips
    gr.HTML("""
        <div class="arch-badge">
            <span class="chip chip-foundry">☁️ Azure AI Foundry</span>
            <span class="chip chip-fabric">🏭 Microsoft Fabric</span>
            <span class="chip chip-ml">🧠 3 ML Models</span>
            <span class="chip chip-model">⚡ GPT-5.1</span>
        </div>
    """)

    # Welcome card
    gr.HTML("""
        <div class="welcome-card">
            <span class="wave">👋</span>
            <h2>Hi, I'm Agent Ken — your AI Companion for Product Management & Data</h2>
            <p>I help you learn product management, explore AI technology, and discover insights from
            my hands-on experience with Microsoft Fabric — including data pipelines, ML models,
            and turning raw data into real product decisions.</p>
            <div class="welcome-tags">
                <span>📊 Live Product Data</span>
                <span>🔮 Churn Prediction</span>
                <span>💰 LTV Analysis</span>
                <span>⚠️ Anomaly Detection</span>
                <span>🧪 A/B Test Results</span>
                <span>📝 PRDs & Strategy</span>
                <span>🤖 AI & Tech</span>
            </div>
        </div>
    """)

    # Data overview cards
    gr.HTML("""
        <div class="data-section">
            <h3>📊 Connected Product Data (Fabric → Foundry)</h3>
            <div class="data-grid">
                <div class="data-card">
                    <span class="number">5,000</span>
                    <span class="label">Users Tracked</span>
                </div>
                <div class="data-card">
                    <span class="number">204K</span>
                    <span class="label">Events Analyzed</span>
                </div>
                <div class="data-card">
                    <span class="number">90 Days</span>
                    <span class="label">Data Window</span>
                </div>
                <div class="data-card">
                    <span class="number">25.6%</span>
                    <span class="label">D7 Retention</span>
                </div>
                <div class="data-card">
                    <span class="number">$19.87</span>
                    <span class="label">Avg LTV</span>
                </div>
                <div class="data-card">
                    <span class="number">5</span>
                    <span class="label">A/B Tests Run</span>
                </div>
            </div>
        </div>
    """)

    # Upload hint
    gr.HTML("""
        <div class="upload-hint">
            📎 <strong>You can also upload your files!</strong> (Supported: CSV, Excel, PDF, Word, Images, TXT, JSON.)
            I'll analyze them and give you insights — and if you want, I can compare your data against our Fabric benchmark.
        </div>
    """)

    # Chat
    chatbot = gr.Chatbot(label="Agent Ken", height=440)
    thread_state = gr.State(value=None)

    # Input row
    with gr.Row():
        msg_input = gr.Textbox(
            placeholder="Ask about metrics, upload a file, or any PM/tech question...",
            label="", show_label=False, scale=5, lines=1, max_lines=3, container=False,
        )
        send_btn = gr.Button("Send 🚀", variant="primary", scale=1, min_width=100)

    with gr.Row():
        clear_btn = gr.Button("🗑️ New Conversation", variant="secondary", size="sm", scale=1)

    # File upload
    with gr.Accordion("📎 Upload a file (CSV, Excel, PDF, Word, Images, JSON)", open=False):
        file_input = gr.File(
            label="",
            file_count="multiple",
            file_types=[
                ".csv", ".xlsx", ".xls", ".json",
                ".txt", ".md", ".pdf",
                ".doc", ".docx", ".pptx",
                ".png", ".jpg", ".jpeg", ".gif",
                ".py", ".html", ".xml",
            ],
        )

    with gr.Accordion("💡 Try an example", open=False):
        gr.Examples(examples=EXAMPLES, inputs=msg_input, label="")

    # Footer
    gr.HTML("""
        <div class="footer-info">
            Built by <strong>Kendrick Filbert</strong><br>
            Azure AI Foundry (GPT-5.1) · Microsoft Fabric (OneLake + MLflow) · 3 ML Models (Churn · LTV · Anomaly)<br>
            Supports: CSV · Excel · PDF · Word · Images · JSON · TXT
        </div>
    """)

    # Events
    send_btn.click(
        fn=respond,
        inputs=[msg_input, file_input, chatbot, thread_state],
        outputs=[chatbot, thread_state],
    ).then(lambda: ("", None), outputs=[msg_input, file_input])

    msg_input.submit(
        fn=respond,
        inputs=[msg_input, file_input, chatbot, thread_state],
        outputs=[chatbot, thread_state],
    ).then(lambda: ("", None), outputs=[msg_input, file_input])

    clear_btn.click(fn=new_conversation, outputs=[chatbot, thread_state])

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        css=CSS,
        theme=gr.themes.Soft(),
        ssr_mode=False,
    )