import json
import pandas as pd
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

ORDER_COLUMNS = [
    "customer",
    "item",
    "quantity",
    "flavor",
    "pickup_time",
    "delivery_address",
    "payment_status",
    "notes",
    "missing_fields",
]

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model.eval()

# SYSTEM_PROMPT = """
# You extract customer orders from messy DMs for tiny sellers.
# Return only valid JSON with this exact shape:
# {
#   "orders": [
#     {
#       "customer": "",
#       "item": "",
#       "quantity": "",
#       "flavor": "",
#       "pickup_time": "",
#       "delivery_address": "",
#       "payment_status": "",
#       "notes": "",
#       "missing_fields": []
#     }
#   ],
#   "prep_list": [],
#   "reply_drafts": [
#     {"customer": "", "reply": ""}
#   ]
# }
# Use empty strings for unknown values. Put missing details in missing_fields.
# """

# EXAMPLE_INPUT = """Maya: Hi! Can I get 2 dozen cupcakes for Saturday morning? Half vanilla, half chocolate.
# Sam: Need 1 birthday cake, chocolate, for pickup Friday 5pm. I can pay Venmo.
# Lena: Do you still have lemon bars? I need some for tomorrow but not sure how many yet.
# Chris: 12 cookies please, pickup at the farmers market. Paid already.
# """

SYSTEM_PROMPT = """
You are a careful order extraction engine for tiny sellers.

Extract customer orders from messy DMs. Return only valid JSON with this exact shape:
{
  "orders": [
    {
      "customer": "",
      "item": "",
      "quantity": "",
      "flavor": "",
      "pickup_time": "",
      "delivery_address": "",
      "payment_status": "",
      "notes": "",
      "missing_fields": []
    }
  ],
  "prep_list": [],
  "reply_drafts": [
    {"customer": "", "reply": ""}
  ]
}

Critical rules:
- Do not invent customer names.
- Do not change customer names.
- Do not merge different customers.
- Include every customer message that looks like an order or possible order.
- Use only facts explicitly present in the messages.
- If a value is unknown, use an empty string.
- Do not add order_id or total_cost unless the message mentions them.
- For pickup orders, put the pickup time in pickup_time and leave delivery_address empty unless a real address is given.
- If the customer is unsure, still include the order and put uncertainty in notes.
- missing_fields should only include practical fields the seller needs to ask for, such as quantity, flavor, pickup_time, delivery_address, or payment_status.
- Reply drafts should be short, friendly, and ask only for missing information.
"""

def extract_json(text):
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1:
        raise ValueError("No JSON object found")
    return json.loads(text[start:end + 1])

def normalize_orders(data):
    rows = []
    for order in data.get("orders", []):
        row = {}
        for col in ORDER_COLUMNS:
            value = order.get(col, "")
            if isinstance(value, list):
                value = ", ".join(str(v) for v in value)
            row[col] = value
        rows.append(row)
    return pd.DataFrame(rows, columns=ORDER_COLUMNS)

def format_list(title, items):
    if not items:
        return f"### {title}\nNothing found."
    lines = []
    for item in items:
        if isinstance(item, dict):
            lines.append("- " + json.dumps(item, ensure_ascii=False))
        else:
            lines.append(f"- {item}")
    return f"### {title}\n" + "\n".join(lines)

def format_replies(replies):
    if not replies:
        return "### Reply drafts\nNothing found."
    lines = []
    for reply in replies:
        customer = reply.get("customer", "Customer")
        text = reply.get("reply", "")
        lines.append(f"**{customer}**\n\n{text}")
    return "### Reply drafts\n\n" + "\n\n---\n\n".join(lines)

def analyze_messages(messages):
    if not messages.strip():
        return pd.DataFrame(columns=ORDER_COLUMNS), "Paste some DMs first.", "", ""

    prompt = tokenizer.apply_chat_template(
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": messages},
        ],
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=900,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )

    try:
        data = extract_json(generated)
    except Exception as exc:
        return (
            pd.DataFrame(columns=ORDER_COLUMNS),
            f"### Needs review\nThe model did not return valid JSON: {exc}",
            "",
            generated,
        )

    orders_df = normalize_orders(data)
    prep = format_list("Prep list", data.get("prep_list", []))
    replies = format_replies(data.get("reply_drafts", []))
    raw = json.dumps(data, indent=2, ensure_ascii=False)
    return orders_df, prep, replies, raw

with gr.Blocks(title="DM Order Desk") as demo:
    gr.Markdown("# DM Order Desk")
    gr.Markdown("Turn messy customer DMs into clean orders, prep lists, and reply drafts using a small model.")

    with gr.Row():
        with gr.Column(scale=1):
            messages = gr.Textbox(
                label="Messy customer DMs",
                value=EXAMPLE_INPUT,
                lines=14,
            )
            run = gr.Button("Organize orders", variant="primary")

        with gr.Column(scale=2):
            orders = gr.Dataframe(label="Order sheet", headers=ORDER_COLUMNS)
            prep = gr.Markdown(label="Prep list")
            replies = gr.Markdown(label="Reply drafts")
            raw = gr.Code(label="Raw JSON", language="json")

    run.click(analyze_messages, inputs=messages, outputs=[orders, prep, replies, raw])

demo.launch()