"""
OCR Service — extracts transaction data from receipt images using Groq vision.

Model: meta-llama/llama-4-scout-17b-16e-instruct (multimodal)
Input: image file path or bytes
Output: dict with amount, date, merchant, category_suggestion, notes
"""

import base64
import json
import logging
import os
import re
from datetime import date, datetime
from flask import current_app

log = logging.getLogger(__name__)

GROQ_VISION_MODEL = 'meta-llama/llama-4-scout-17b-16e-instruct'
SUPPORTED_MIME_TYPES = {
    '.jpg':  'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.png':  'image/png',
    '.gif':  'image/gif',
    '.webp': 'image/webp',
}

RECEIPT_PROMPT = """You are a receipt data extractor. Analyze this receipt image and extract the following information.

Respond ONLY with a valid JSON object, no explanation, no markdown, no code fences. Just raw JSON.

Required format:
{
  "amount": <total amount as a number, no currency symbols, e.g. 85.50>,
  "date": "<date in YYYY-MM-DD format, or null if not found>",
  "merchant": "<store/restaurant/vendor name, or null if not found>",
  "category_suggestion": "<one of: Food & Dining, Transport, Shopping, Health, Entertainment, Utilities, Housing, Education, Personal Care, Travel, Other>",
  "notes": "<brief description, e.g. 'Lunch at McDonald\\'s' or 'Grocery shopping'>"
}

Rules:
- amount: use the TOTAL/GRAND TOTAL amount, not subtotal. Numbers only, no commas.
- date: format as YYYY-MM-DD. If only month/day visible, use current year.
- merchant: the business name only, no address.
- category_suggestion: pick the single most appropriate category from the list.
- notes: keep under 60 characters.
- If you cannot read the receipt clearly, still return your best guess with available data.
- Never return null for amount if any number is visible."""


def extract_from_file(file_path):
    """
    Extract receipt data from an image file on disk.
    Returns dict or None on failure.
    """
    ext = os.path.splitext(file_path)[1].lower()
    mime_type = SUPPORTED_MIME_TYPES.get(ext)

    if not mime_type:
        log.warning(f'[ocr] Unsupported file type: {ext}')
        return _error_result(f'Unsupported file type: {ext}')

    try:
        with open(file_path, 'rb') as f:
            image_bytes = f.read()
        return extract_from_bytes(image_bytes, mime_type)
    except FileNotFoundError:
        log.error(f'[ocr] File not found: {file_path}')
        return _error_result('Receipt file not found')
    except Exception as e:
        log.error(f'[ocr] File read error: {e}')
        return _error_result(str(e))


def extract_from_bytes(image_bytes, mime_type='image/jpeg'):
    """
    Extract receipt data from raw image bytes.
    Returns dict with keys: amount, date, merchant, category_suggestion, notes, error
    """
    api_key = current_app.config.get('GROQ_API_KEY', '')
    if not api_key:
        return _error_result('GROQ_API_KEY not configured')

    # Encode to base64
    b64 = base64.b64encode(image_bytes).decode('utf-8')
    data_url = f'data:{mime_type};base64,{b64}'

    try:
        import requests
        resp = requests.post(
            'https://api.groq.com/openai/v1/chat/completions',
            headers={
                'Authorization': f'Bearer {api_key}',
                'Content-Type': 'application/json',
            },
            json={
                'model': GROQ_VISION_MODEL,
                'messages': [
                    {
                        'role': 'user',
                        'content': [
                            {
                                'type': 'image_url',
                                'image_url': {'url': data_url},
                            },
                            {
                                'type': 'text',
                                'text': RECEIPT_PROMPT,
                            },
                        ],
                    }
                ],
                'max_tokens': 512,
                'temperature': 0.1,  # low temp for consistent structured output
            },
            timeout=30,
        )
        resp.raise_for_status()
        content = resp.json()['choices'][0]['message']['content'].strip()
        log.info(f'[ocr] Raw response: {content[:200]}')
        return _parse_response(content)

    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 400:
            # Model may not support this image — try to parse error
            try:
                err_body = e.response.json()
                msg = err_body.get('error', {}).get('message', str(e))
            except Exception:
                msg = str(e)
            log.error(f'[ocr] Groq 400 error: {msg}')
            return _error_result(f'Image processing failed: {msg[:100]}')
        elif e.response.status_code == 429:
            return _error_result('Rate limit reached. Try again shortly.')
        elif e.response.status_code == 401:
            return _error_result('Invalid Groq API key.')
        else:
            log.error(f'[ocr] HTTP {e.response.status_code}: {e}')
            return _error_result(f'API error ({e.response.status_code})')
    except Exception as e:
        log.error(f'[ocr] Unexpected error: {e}')
        return _error_result('OCR service temporarily unavailable')


def _parse_response(content):
    """Parse the JSON response from Groq into a clean dict."""
    # Strip any accidental markdown fences
    content = re.sub(r'^```(?:json)?\s*', '', content, flags=re.MULTILINE)
    content = re.sub(r'\s*```$', '', content, flags=re.MULTILINE)
    content = content.strip()

    try:
        data = json.loads(content)
    except json.JSONDecodeError:
        # Try to extract JSON object from surrounding text
        match = re.search(r'\{.*\}', content, re.DOTALL)
        if match:
            try:
                data = json.loads(match.group())
            except json.JSONDecodeError:
                log.error(f'[ocr] Could not parse JSON: {content[:200]}')
                return _error_result('Could not parse OCR response')
        else:
            return _error_result('No JSON found in OCR response')

    # Normalise amount
    amount = data.get('amount')
    if amount is not None:
        try:
            amount = float(str(amount).replace(',', '').strip())
            if amount <= 0:
                amount = None
        except (ValueError, TypeError):
            amount = None

    # Normalise date
    raw_date = data.get('date')
    parsed_date = None
    if raw_date:
        for fmt in ('%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%Y/%m/%d'):
            try:
                parsed_date = datetime.strptime(str(raw_date).strip(), fmt).date().isoformat()
                break
            except ValueError:
                continue

    # Map category suggestion to our system category names
    cat_map = {
        'food & dining': 'Food & Dining',
        'food':          'Food & Dining',
        'dining':        'Food & Dining',
        'transport':     'Transport',
        'transportation':'Transport',
        'shopping':      'Shopping',
        'health':        'Health',
        'entertainment': 'Entertainment',
        'utilities':     'Utilities',
        'housing':       'Housing',
        'education':     'Education',
        'personal care': 'Personal Care',
        'travel':        'Travel',
        'other':         'Other',
    }
    raw_cat = str(data.get('category_suggestion', '')).lower().strip()
    category = cat_map.get(raw_cat, data.get('category_suggestion', 'Other'))

    return {
        'amount':              amount,
        'date':                parsed_date or date.today().isoformat(),
        'merchant':            data.get('merchant') or '',
        'category_suggestion': category,
        'notes':               data.get('notes') or '',
        'error':               None,
        'raw':                 data,
    }


def _error_result(msg):
    return {
        'amount':              None,
        'date':                date.today().isoformat(),
        'merchant':            '',
        'category_suggestion': 'Other',
        'notes':               '',
        'error':               msg,
        'raw':                 {},
    }