import os
import gradio as gr
import requests
import json
from datetime import datetime

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_COMPLETION_URL = "https://api.groq.com/openai/v1/chat/completions"
GROQ_MODEL = "llama-3.1-8b-instant"

classification_history = []

def groq_completion(prompt, sys_prompt=None):
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    body = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": sys_prompt or "You are a fast and reliable business email assistant."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3,
        "max_tokens": 512
    }
    try:
        response = requests.post(GROQ_COMPLETION_URL, headers=headers, json=body)
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]
    except Exception:
        return "ERROR:model_request"

def email_classifier_router(raw_email):
    prompt = (
        "Your reply MUST BE valid compact JSON. NO TEXT OR EXPLANATION before or after the JSON. "
        "Given the business email below, extract:"
        "\n- Category (Support Request, Sales Inquiry, Finance/Billing, Urgent Incident, Spam/Marketing)"
        "\n- Priority (Low, Medium, High)"
        "\n- Suggested Recipient"
        "\n- Draft Response (professional reply text)"
        "\n- Summary (one sentence)"
        "\n- Action Items (numbered list)\n"
        "Format: {\"Category\":..., \"Priority\":..., \"Suggested Recipient\":..., \"Draft Response\":..., \"Summary\":..., \"Action Items\":[...]}\n\n"
        f"{raw_email}"
    )
    raw_result = groq_completion(prompt)
    if raw_result == "ERROR:model_request":
        return {
            "Category": "ERROR",
            "Priority": "ERROR",
            "Suggested Recipient": "ERROR",
            "Draft Response": "Model request failed, please check API or retry.",
            "Summary": "ERROR",
            "Action Items": ["ERROR"]
        }
    try:
        output = json.loads(raw_result)
        for key in ["Category", "Priority", "Suggested Recipient", "Draft Response", "Summary", "Action Items"]:
            if key not in output:
                output[key] = ""
    except Exception:
        output = {
            "Category": "PARSE ERROR",
            "Priority": "PARSE ERROR",
            "Suggested Recipient": "",
            "Draft Response": f"Could not extract valid JSON. Raw LLM output:\n{raw_result}",
            "Summary": "PARSE ERROR",
            "Action Items": ["PARSE ERROR"]
        }
    return output

def add_to_history(email, cat, pri, summ, dra):
    entry = {
        "datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "email": email,
        "category": cat,
        "priority": pri,
        "summary": summ,
        "draft": dra
    }
    classification_history.append(entry)

css = """
body { background: linear-gradient(120deg,#10193a 0%,#175ad7 120%)!important; font-family:'Inter','Segoe UI',Arial,sans-serif;}
.gradio-container { background: transparent !important; }
.main-pill-header {
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    margin: 28px auto 23px auto;
    width: fit-content;
    padding: 0;
}
.main-pill {
    background:linear-gradient(90deg,#3084f4 0%,#39d5ff 97%);
    border-radius:34px;
    padding:33px 66px 30px 60px;
    box-shadow:0 0 90px 20px #42e5ff3a,0 0 0 12px #11b6ff0a;
    display:flex;flex-direction:column;align-items:center;
    border:3.5px solid #79e8ffcc;
}
.main-pill-title {
    font-size:3.2em;font-weight:900;color:#fff;letter-spacing:.04em;
    text-shadow:0 10px 70px #1ad1f3a9,0 3px 24px #36b0ecba;
}
.main-pill-subtitle {
    color:#e7f7ff;margin-top:11px;font-size:1.17em;font-weight:500;letter-spacing:.01em;text-shadow:0 2px 14px #99faffad;
}
.section-header {
    font-size:1.55em; font-weight:800; letter-spacing:.02em; color:#fff;
    margin-top:28px;
    margin-bottom:0px !important; /* No gap below header */
    text-shadow:0 4px 34px #09dede,0 2px 8px #2989d2; display:block;
}
.card-block {
    background:rgba(27,36,68,0.96); border-radius:30px; box-shadow:0 6px 42px #0b2269bb,0 1.5px 0 #186fc055 inset;
    padding:38px 28px 29px 28px; margin:0 0 32px 0 !important; color:#fafeff; border:2px solid #2847a46c; backdrop-filter:blur(8px);
}
.card-gap { height: 5px; }
.result-labels-row{display:flex;align-items:center;gap:18px;margin-bottom:10px;margin-top:-6px;}
.result-label-text{font-size:1.08em;font-weight:700;color:#daf6ff;letter-spacing:.01em;margin-right:2px;}
.badge{display:inline-block;border-radius:999px;padding:4px 16px;font-weight:600;background:linear-gradient(90deg,#2870ee 70%,#13e4fc 140%);box-shadow:0 1px 12px #0663fd29;color:#fff;margin-right:8px;font-size:1.02em;vertical-align:middle;}
.priority-high{background:linear-gradient(90deg,#ed4864 60%,#fa4764 150%)!important;color:#fff;}
.priority-medium{background:linear-gradient(90deg,#f3aa36 70%,#f6ea48 120%)!important;color:#222;}
.priority-low{background:linear-gradient(90deg,#18cd67 60%,#27f8ac 150%)!important;color:#fff;}
.history-cardrow {display:flex;flex-wrap:wrap;gap:30px;}
.history-record {background:rgba(35,54,90,0.92);color:#e2f5ff;margin:0 0 13px 0; border-radius:16px; border-left:5px solid #10baff; box-shadow:0 2px 12px #08526733;padding:20px 30px 16px 30px;min-width:340px;max-width:380px;box-sizing:border-box;}
.history-top { display:flex; align-items:center; justify-content:space-between;}
.history-catlab { font-size:1em;font-weight:600;margin-top:3px; }
.history-summ { font-style:italic;font-size:1em;margin-top:10px;color:#eaf7ff;}
.history-btnrow { display:flex; gap:12px; margin-top:16px;}
.btn-md {padding:10px 22px;border-radius:14px; border:none;font-size:1.03em;font-weight:700;background:#3246b8; color:#fff;}
.btn-md:hover { background:#2284cc;}
.clear-btn { background:#db2828!important;color:#fff!important;padding:10px 24px;font-weight:600;border-radius:13px;border:none;font-size:1.07em;margin-bottom:22px;}
#paste-btn .wrap {font-size:1.13em;}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("""
    <div class='main-pill-header'>
      <div class='main-pill'>
        <span class='main-pill-title'>📧 AI Email Classifier & Router</span>
        <span class='main-pill-subtitle'>Instantly categorize, summarize, and draft replies to business emails.</span>
      </div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=6, min_width=420):
            gr.HTML("<div class='section-header'>✉️ Paste Your Email</div>")
            gr.HTML("<div class='card-gap'></div>")
            with gr.Group(elem_id="input-card", elem_classes=["card-block"]):
                gr.Markdown("<span style='font-size:1.08em; color:#a6c5e7;'>Choose an example:</span>")
                with gr.Row():
                    support_button      = gr.Button("Support Request")
                    sales_button        = gr.Button("Sales Inquiry")
                    finance_button      = gr.Button("Finance / Billing")
                    urgent_button       = gr.Button("Urgent Incident")
                    spam_button         = gr.Button("Spam / Marketing")
                    account_button      = gr.Button("Account Issue")
                    legal_button        = gr.Button("Legal Inquiry")
                    meeting_button      = gr.Button("Meeting Scheduling")
                email_box = gr.Textbox(lines=8, label="", placeholder="From: ...\nSubject: ...")
                paste_btn = gr.Button("📋 Paste from Clipboard", elem_id="paste-btn", elem_classes=["glass-btn"])
                classify_btn = gr.Button("Classify Email", elem_classes=["glass-btn"])
        with gr.Column(scale=8, min_width=430):
            gr.HTML("<div class='section-header'>🧮 Classification Results</div>")
            gr.HTML("<div class='card-gap'></div>")
            with gr.Group(elem_id="output-card", elem_classes=["card-block"]):
                result_labels = gr.HTML()
                recipient = gr.Textbox(label="Suggested Recipient")
                draft_response = gr.Textbox(lines=5, label="Generated Response Draft")
                summary = gr.Textbox(label="Summary")
                action_items = gr.Textbox(lines=3, label="Extracted Action Items")

    gr.HTML("<div class='section-header'>🕒 Classification History</div>")
    gr.HTML("<div class='card-gap'></div>")
    with gr.Group(elem_id="history-card", elem_classes=["card-block"]):
        clear_btn = gr.Button("Clear History", elem_classes=["clear-btn"])
        history_cardbox = gr.HTML(elem_id="history-cardbox")

    examples = [
        "From: support@mycrm.com\nSubject: Cannot login to dashboard\nBody: I've been unable to sign into my account for 3 days even after password reset. Tried from multiple browsers and mobile apps.",
        "From: prospect@b2bsales.com\nSubject: Inquiry about pricing\nHi, Can you send over your enterprise pricing, feature list, contract terms, and onboarding timeline for Q1 2025?",
        "From: accounts@acmecorp.com\nSubject: Invoice #00456 Payment Reminder\nBody: This is a reminder that your invoice #00456 for $4,200.00 is due next week. Please remit payment online or contact billing for assistance.",
        "From: jane.doe@company.org\nSubject: Meeting request\nCould you schedule a call with the marketing team for Friday 3pm? Please include John and Lisa from design.",
        "From: admin@infra.com\nSubject: URGENT Server Down\nThe production DB went down at 11:22 UTC, please escalate to engineering immediately. All web portals affected.",
        "From: legal@mycrm.com\nSubject: Contract question\nCan your team confirm paragraph 7.4 in the service agreement applies to all regional resellers? Our counsel needs official confirmation.",
        "From: billing@b2bsales.com\nSubject: Account update\nPlease update billing address for account #AC10239 to 44 Queen St, Central, Hong Kong. Email invoice thereafter.",
        "From: teamlead@company.org\nSubject: Follow-up on project\nHi, just wanted a final status update on our July roadmap deliverables. Are we on track to ship by the end of month?",
    ]
    support_button.click(lambda: examples[0], None, email_box)
    sales_button.click(lambda: examples[1], None, email_box)
    finance_button.click(lambda: examples[2], None, email_box)
    urgent_button.click(lambda: examples[3], None, email_box)
    spam_button.click(lambda: examples[4], None, email_box)
    account_button.click(lambda: examples[5], None, email_box)
    legal_button.click(lambda: examples[6], None, email_box)
    meeting_button.click(lambda: examples[7], None, email_box)

    # Clipboard paste to input box!
    paste_btn.click(None, None, email_box, js="(inputs,outputs)=>navigator.clipboard.readText()")

    def classify_and_render(email_text):
        result = email_classifier_router(email_text)
        cat = result.get("Category", "")
        pri = result.get("Priority", "")
        rec = result.get("Suggested Recipient", "")
        dra = result.get("Draft Response", "")
        summ = result.get("Summary", "")
        acts = result.get("Action Items", [])
        acts = "\n".join(acts) if isinstance(acts, list) else acts
        add_to_history(email_text, cat, pri, summ, dra)
        badge_html = (
            "<div class='result-labels-row'>"
            "<span class='result-label-text'>Category:</span>"
            f"<span class='badge'>{cat}</span>"
            "<span class='result-label-text'>Priority:</span>"
            f"<span class='badge priority-{pri.lower()}'>{pri}</span>"
            "</div>"
        )
        # Only show last 5 records
        recent_history = classification_history[-5:][::-1]
        history_html = "<div class='history-cardrow'>"
        for idx, h in enumerate(recent_history):
            history_html += f"""
            <div class='history-record'>
                <div class='history-top'>
                    <span class='history-ts'>{h['datetime']}</span>
                    <span><span class='badge'>{h['category']}</span>
                    <span class='badge priority-{h['priority'].lower()}'>{h['priority']}</span></span>
                </div>
                <div class='history-summ'>{h['summary']}</div>
                <div class='history-btnrow'>
                    <button class='btn-md' onclick="navigator.clipboard.writeText({json.dumps(h['email'])});">Copy Email</button>
                    <button class='btn-md' onclick="window.parent.postMessage({json.dumps(h['email'])}, '*');">Load</button>
                </div>
            </div>
            """
        history_html += "</div>"
        return badge_html, rec, dra, summ, acts, history_html

    def clear_history():
        classification_history.clear()
        return "<div class='history-cardrow'></div>"

    classify_btn.click(
        classify_and_render,
        inputs=[email_box],
        outputs=[result_labels, recipient, draft_response, summary, action_items, history_cardbox]
    )
    clear_btn.click(clear_history, None, history_cardbox)

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