import spaces
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import json
import re

print("Loading Qwen2.5-1.5B-Instruct...")
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.float16, device_map="auto")
print("Ready!")

@spaces.GPU
def extract(email):
    if not email.strip():
        return "Paste an email first"
    
    messages = [
        {"role": "system", "content": "You extract obligations from client emails. Respond with ONLY valid JSON, no other text."},
        {"role": "user", "content": f"""Analyze this client email and extract obligations.

Email:
{email}

Respond with ONLY this JSON format:
{{
  "client": "client name and company",
  "project": "what they want built",
  "promises": ["deliverable 1", "deliverable 2"],
  "owes": ["what client must provide"],
  "deadline": "when due",
  "payment": "amount and terms",
  "red_flags": ["scope creep or vague items"]
}}"""}
    ]
    
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=500,
        temperature=0.2,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    
    match = re.search(r'\{.*\}', response, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group())
            out = f"### 📋 {data.get('client', '?')}\n\n"
            out += f"**Project:** {data.get('project', '?')}\n\n"
            out += f"**📅 Deadline:** {data.get('deadline', '?')}\n\n"
            out += f"**💰 Payment:** {data.get('payment', '?')}\n\n"
            out += "**✅ Dev Promises:**\n"
            for p in data.get('promises', []):
                out += f"- {p}\n"
            out += "\n**🔄 Client Owes:**\n"
            for o in data.get('owes', []):
                out += f"- {o}\n"
            out += "\n**🚩 Red Flags:**\n"
            for r in data.get('red_flags', []):
                out += f"- {r}\n"
            return out
        except Exception as e:
            return f"Parse error: {e}\n\nRaw:\n{response[:500]}"
    return f"No JSON found. Output:\n{response[:500]}"

with gr.Blocks() as app:
    gr.Markdown("# 💼 Obligation Extractor\n**Build Small Hackathon — Qwen2.5 1.5B**\n\nPaste a client email to extract obligations, deadlines, payment terms, and red flags.")
    email = gr.Textbox(label="Paste Client Email", lines=12, placeholder="Paste the full email here...")
    btn = gr.Button("🔍 Extract Obligations", variant="primary", size="lg")
    output = gr.Markdown()
    btn.click(extract, email, output)
    gr.Examples([
        "Hi! Website redesign needed. Budget $5000 (50% upfront, 50% on completion). Due March 15. We provide content and feedback within 24h.",
        "Hi, need a landing page + email signup. Budget around $1500. Pay on completion. Need it ASAP. Can you also do social media graphics?"
    ], email)

app.launch()