import ollama
import json
import base64

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 = """
    You are an expert data extraction assistant. 
    Extract the following fields from the provided document and return them as a SINGLE VALID JSON OBJECT:
    - invoice_number (string)
    - date (string)
    - vendor_name (string)
    - total_amount (string)
    - currency (string)
    - line_items (list of objects with: description, quantity, unit_price, total)

    IMPORTANT: 
    - Return ONLY the JSON. No markdown formatting, no explanations. 
    - If a field is not found, 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 ""}