# app.py
import os
import gradio as gr
import json
import textwrap
import re
import torch

# ZeroGPU decorator helper (provided by the Spaces runtime)
# Note: 'spaces' module is available in the Spaces environment for ZeroGPU usage.
import spaces

from model_loader import load_model_cpu, move_model_to_gpu, move_model_to_cpu, generate_with_model
from db import fetch_client_info

# Load model on CPU at startup (low memory footprint)
MODEL_INFO = load_model_cpu()  # returns dict { "tokenizer":..., "model":..., "base_model_name": ... }

# utility to extract policy or client id
def extract_identifier(email_text: str):
    m = re.search(r"(POL\d+|CUST\d+)", email_text, flags=re.IGNORECASE)
    return m.group(1).upper() if m else None

def format_safe_reply(identifier: str):
    return {
        "type": "unknown",
        "email_subject": "Policy / Client not found",
        "email_body": textwrap.dedent(f"""\
            Dear Customer,

            We could not locate any records for the policy number or client ID mentioned: {identifier}.
            Please verify the details and share the correct policy/client ID so we can assist you further.

            Regards,
            LIC Customer Service Team
            """).strip(),
        "raw_output": None
    }

def clean_and_package(raw_text: str, email_type: str, client_data: dict):
    # clean duplicates and stray markers
    resp = raw_text.split("### Response:")[-1]
    resp = resp.split("### Output:")[0]
    resp = re.sub(r"\n{3,}", "\n\n", resp).strip()

    subject_map = {
        "query": "Policy Information Inquiry",
        "complaint": "Issue Resolution Update",
        "request": "Service Request Confirmation",
        "feedback": "Thank You for Your Feedback"
    }
    subject = subject_map.get(email_type.lower(), "Response from LIC Customer Service")
    name = client_data.get("name") if client_data else "Valued Customer"

    body = textwrap.dedent(f"""\
        Dear {name},

        {resp}

        Regards,
        LIC Customer Service Team
    """).strip()

    return {
        "type": email_type.lower(),
        "email_subject": subject,
        "email_body": body,
        "raw_output": raw_text.strip()
    }

# GPU-only inference: allocate GPU only during this call
@spaces.GPU(duration=45)   # seconds; adjust if needed
def infer_and_respond(email_text: str, mail_type: str):
    # find identifier and fetch client info
    identifier = extract_identifier(email_text)
    client = fetch_client_info(identifier) if identifier else None

    if not client:
        return format_safe_reply(identifier or "None provided")

    # Move model to GPU, generate, then move back to CPU
    move_model_to_gpu(MODEL_INFO)   # will .to('cuda') and empty cache as needed
    try:
        raw = generate_with_model(MODEL_INFO, email_text, mail_type, client)
    finally:
        # Ensure we free GPU quickly
        move_model_to_cpu(MODEL_INFO)

    return clean_and_package(raw, mail_type, client)

def app_handle(email_text):    # frontend wrapper - keep syncronous
        # Let the model classify the email type automatically
    mail_type = "Query"  # default, model will determine actual type
    result = infer_and_respond(email_text, mail_type)
    formatted_text = format_response_text(result)
    return formatted_text

def format_response_text(result):
    """Format JSON response into a beautiful email-like preview."""
    if not result:
        return "No response generated."

    email_type = result.get("type", "unknown").title()
    subject = result.get("email_subject", "No Subject")
    body = result.get("email_body", "")
    raw = result.get("raw_output", "")

    formatted = f"""
<div style="border:1px solid #e5e5e5; border-radius:10px; padding:20px; background:#fafafa; font-family:Arial, sans-serif;">
    
<h3 style="margin-top:0; color:#2c3e50;">📧 LIC Email Response</h3>

<p><strong>Type:</strong> {email_type}</p>

<p><strong>Subject:</strong><br>
<span style="color:#34495e; font-size:16px;">{subject}</span></p>

<hr style="border:0; border-top:1px solid #ddd; margin:15px 0;">

<p style="white-space:pre-wrap; line-height:1.6; color:#333;">
{body}
</p>

<hr style="border:0; border-top:1px solid #ddd; margin:15px 0;">

<p style="font-size:13px; color:#7f8c8d;">
Generated by <strong>Insurance CRM Assistant</strong>  
This is an automated response based on client records.
</p>

</div>
"""

    return formatted



title = "Insurance CRM Email Assistant"
description = "Paste a customer email; the assistant classifies it, fetches relevant client data (by policy/client id), and returns a JSON response."

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}\n\n{description}")
    with gr.Row():
        inp = gr.Textbox(lines=6, placeholder="Paste the customer email here...", label="Customer Email")
    formatted_out = gr.Markdown(label="Formatted Email Response")
    btn = gr.Button("Generate Response")
    btn.click(fn=app_handle, inputs=[inp], outputs=[formatted_out])
if __name__ == "__main__":
    demo.launch()