import ollama
import json
import base64

INVOICE_SCHEMA = {
    "document_type": None,
    "invoice_number": None,
    "invoice_date": None,
    "due_date": None,
    "purchase_order_number": None,
    "vendor": {
        "name": None,
        "address": None,
        "email": None,
        "phone": None,
        "gstin": None,
        "tax_id": None,
        "website": None
    },
    "customer": {
        "name": None,
        "address": None,
        "gstin": None
    },
    "amounts": {
        "subtotal": None,
        "tax": None,
        "discount": None,
        "shipping": None,
        "round_off": None,
        "total": None,
        "amount_paid": None,
        "balance_due": None,
        "currency": None
    },
    "tax_breakdown": [
        {
            "type": None,
            "rate": None,
            "amount": None
        }
    ],
    "line_items": [
        {
            "line_no": None,
            "description": None,
            "product_code": None,
            "hsn_sac": None,
            "quantity": None,
            "unit": None,
            "unit_price": None,
            "discount": None,
            "tax_rate": None,
            "tax_amount": None,
            "total": None
        }
    ],
    "payment_information": {
        "bank_name": None,
        "account_number": None,
        "ifsc": None,
        "upi_id": None
    },
    "metadata": {
        "pages": None,
        "ocr_confidence": None,
        "language": None
    }
}

def extract_data(text: str = None, image_path: str = None, model_type: str = "text") -> dict:
    """
    Extracts structured data using either Text (Gemma) or Vision (Qwen) models.
    """
    
    prompt = f"""
    You are an expert data extraction assistant. 
    Extract every possible detail from the provided document and return it strictly as a SINGLE VALID JSON OBJECT matching the following schema structure:
    
    {json.dumps(INVOICE_SCHEMA, indent=4)}

    IMPORTANT: 
    - Return ONLY the JSON. No markdown formatting, no explanations, no prefix.
    - If a field is not found or data is not available, use null.
    """
    
    messages = [{'role': 'user', 'content': prompt}]
    model = 'gemma:2b'
    
    if model_type == 'vision':
        if not image_path:
            return {"error": "Image path required for vision mode"}
        
        # Qwen-VL handles images passed in the message
        model = 'qwen2.5vl:7b' # Using the installed model ID
        messages[0]['images'] = [image_path]
        messages[0]['content'] = "Analyze this image. " + prompt
    else:
        # Text Mode
        if not text:
            return {"error": "Text required for text mode"}
        messages[0]['content'] += f"\n\n---\n{text}\n---"

    try:
        response = ollama.chat(model=model, messages=messages)
        content = response['message']['content']
        
        # Clean up markdown
        content = content.replace("```json", "").replace("```", "").strip()
        
        return json.loads(content)

    except Exception as e:
        print(f"LLM Extraction Error ({model_type}): {e}")
        return {"error": str(e), "raw_output": content if 'content' in locals() else ""}