import os
from groq import Groq
from dotenv import load_dotenv

load_dotenv()


client = Groq(api_key=os.getenv("GROQ_API_KEY"))

MODEL = "llama-3.3-70b-versatile"


SYSTEM_PROMPT = """
You are a helpful receipt analysis assistant.
The user will give you structured data extracted from a receipt image,
along with a question about it.

Your job:
- Answer the question clearly and concisely
- Use the receipt data provided — do not make up values
- Format currency values clearly
- If the data doesn't contain enough information to answer, say so

Keep answers short and direct. No unnecessary explanation.
"""



def ask_about_receipt(receipt_json: dict, user_question: str) -> str:
    receipt_text = format_receipt_for_llm(receipt_json)

    messages = [
        {
            "role"   : "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role"   : "user",
            "content": f"""
Here is the receipt data:

{receipt_text}

Question: {user_question}
"""
        }
    ]

    response = client.chat.completions.create(
        model      = MODEL,
        messages   = messages,
        temperature= 0.2,    
        max_tokens = 512
    )

    return response.choices[0].message.content.strip()



def format_receipt_for_llm(receipt_json: dict) -> str:
    lines = []
    lines.append(f"Vendor   : {receipt_json.get('vendor', 'unknown')}")
    lines.append("Items    :")

    for item in receipt_json.get("items", []):
        name  = item.get("name",  "")
        count = item.get("count", "")
        price = item.get("price", "")
        lines.append(f"  - {name:<25} {count:<6} {price}")

    lines.append(f"Subtotal : {receipt_json.get('subtotal', 'N/A')}")
    lines.append(f"Tax      : {receipt_json.get('tax',      'N/A')}")
    lines.append(f"Service  : {receipt_json.get('service',  'N/A')}")
    lines.append(f"Total    : {receipt_json.get('total',    'N/A')}")

    return "\n".join(lines)


    sample_receipt = {
        "vendor"  : "unknown",
        "items"   : [
            {"name": "Nasi Campur Bali", "count": "1 x", "price": "75,000"},
            {"name": "Ice Lemon Tea",    "count": "1 x", "price": "24,000"},
            {"name": "MilkShake Strawb", "count": "1 x", "price": "37,000"},
            {"name": "Hot Tea",          "count": "2 x", "price": "44,000"},
        ],
        "subtotal": "1,346,000",
        "tax"     : "144,695",
        "service" : "100,950",
        "total"   : "1,591,600"
    }

    questions = [
        "What is the total amount?",
        "How much did I spend on drinks?",
        "How many items were ordered?",
        "What percentage of the total is tax?"
    ]

    for q in questions:
        print(f"\nQ: {q}")
        print(f"A: {ask_about_receipt(sample_receipt, q)}")