from flask import Flask, request, render_template
from huggingface_hub import InferenceClient
import re

app = Flask(__name__)

# Initialize DeepSeek-R1 client
client = InferenceClient(model="deepseek-ai/deepseek-llm-67b-chat")

def parse_llm_response(response):
    """Improved parsing that handles model's raw responses"""
    result = {
        "Brand": None,
        "Category": None,
        "Gender": None,
        "Price": None
    }
    
    # Enhanced pattern matching for flexible JSON extraction
    patterns = {
        "brand": r'"brand":\s*"([^"]*)"',
        "category": r'"category":\s*"([^"]*)"',
        "gender": r'"gender":\s*"([^"]*)"',
        "price_range": r'"price_range":\s*"([^"]*)"'
    }
    
    for key, pattern in patterns.items():
        match = re.search(pattern, response, re.IGNORECASE)
        if match:
            value = match.group(1).strip()
            if value.lower() in ["null", "n/a", ""]:
                continue
            if key == "brand":
                result["Brand"] = value.title()
            elif key == "category":
                result["Category"] = value.title()
            elif key == "gender":
                result["Gender"] = value.title()
            elif key == "price_range":
                result["Price"] = value.upper()
    
    return result

def analyze_query(query):
    """Enhanced prompt for luxury brand understanding"""
    prompt = f"""Analyze this fashion query and extract structured data. Follow these rules:

1. Brand: Identify the luxury fashion brand mentioned (e.g., Gucci, Prada, Balenciaga)
2. Category: Product type (perfume, bag, shoes, etc.)
3. Gender: men, women, or unisex
4. Price: Exact price range from query

Return JSON format:

{{
    "brand": "<brand name>",
    "category": "<product category>",
    "gender": "<target gender>",
    "price_range": "<price info>"
}}

Query: "{query}"
"""

    response = client.text_generation(
        prompt=prompt,
        max_new_tokens=200,
        temperature=0.01,  # More deterministic output
        stop_sequences=["\n\n"]  # Prevent extra text
    )
    
    return parse_llm_response(response)

@app.route("/", methods=["GET", "POST"])
def index():
    result = None
    query = ""
    if request.method == "POST":
        query = request.form.get("query", "")
        if query.strip():
            result = analyze_query(query)
    return render_template("index.html", result=result, query=query)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)