services/ocr.py

import json
import os
from typing import Any

from google import genai
from google.genai import types

from services.ai_service import gemini_keys, with_gemini_key_fallback


OCR_MODEL = "gemini-2.5-flash"
DEFAULT_OCR_API_KEY = ""

SYSTEM_PROMPT = """
You are a precise business card data extraction assistant.

Extract all available information from the business card image and return ONLY a valid JSON object.

Use the following schema exactly:

{
  "person": {
    "first_name": "",
    "last_name": "",
    "job_title": "",
    "email": "",
    "phone": {
      "mobile": "",
      "office": "",
      "extension": ""
    },
    "social": {
      "linkedin": ""
    }
  },
  "company": {
    "company_name": "",
    "website": "",
    "address": {
      "full_address": "",
      "street": "",
      "city": "",
      "state_province": "",
      "postal_code": "",
      "country": ""
    }
  },
  "metadata": {
    "confidence_score": 0.0,
    "unparsed_text": "",
    "card_language": ""
  }
}

Rules:
1. Return ONLY valid JSON.
2. Do NOT wrap the response in markdown or code fences.
3. Do NOT include explanations or commentary.
4. If a field is missing, use an empty string "".
5. confidence_score must be between 0 and 1.
6. Put any text that cannot be mapped to a field into metadata.unparsed_text.
7. Infer card_language using ISO language codes such as "en", "fr", "de", "es", etc.
8. Extract LinkedIn URLs if present.
9. Separate mobile, office, and extension numbers when possible.
10. Preserve all extracted values exactly as written on the card.
"""


def _get_api_key() -> str:
    keys = gemini_keys()
    return keys[0] if keys else DEFAULT_OCR_API_KEY


def _parse_json_response(text: str) -> dict[str, Any]:
    cleaned = text.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.lower().startswith("json"):
            cleaned = cleaned[4:].strip()
    return json.loads(cleaned)


def extract_business_card(image_bytes: bytes, mime_type: str) -> dict[str, Any]:
    """
    Extract structured business-card data from uploaded image bytes.
    """
    if not image_bytes:
        raise ValueError("No image bytes provided.")

    api_key = _get_api_key()
    if not api_key:
        raise RuntimeError("Set GEMINI_API_KEY_prime_1, GEMINI_API_KEY_prime_2, or GEMINI_API_KEY_prime_3 before using OCR.")

    def call_ocr(key: str):
        client = genai.Client(api_key=key)
        return client.models.generate_content(
            model=OCR_MODEL,
            contents=[
                types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
                SYSTEM_PROMPT,
            ],
            config=types.GenerateContentConfig(service_tier=types.ServiceTier.STANDARD),
        )

    response_text = ""
    try:
        response = with_gemini_key_fallback((), call_ocr)
        response_text = response.text or ""
    except Exception as gemini_err:
        import os
        import base64
        import anthropic
        
        claude_key = os.environ.get("ANTHROPIC_API_KEY", "").strip()
        if not claude_key:
            raise RuntimeError(f"Gemini OCR failed and ANTHROPIC_API_KEY is not configured. Gemini Error: {gemini_err}")
            
        client = anthropic.Anthropic(api_key=claude_key)
        
        image_data = base64.standard_b64encode(image_bytes).decode("utf-8")
        image_block = {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": mime_type,
                "data": image_data,
            },
        }
        
        messages = [
            {
                "role": "user",
                "content": [
                    image_block,
                    {"type": "text", "text": "Extract all available information from the business card image and return ONLY a valid JSON object."},
                ],
            }
        ]

        full_response = ""
        max_continuations = 5
        continuation_count = 0

        while True:
            claude_response = client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=1024,
                system=SYSTEM_PROMPT,
                messages=messages,
            )

            for block in claude_response.content:
                if block.type == "text":
                    full_response += block.text

            stop_reason = claude_response.stop_reason

            if stop_reason == "end_turn":
                break
            elif stop_reason == "max_tokens":
                continuation_count += 1
                if continuation_count >= max_continuations:
                    break
                messages.append({"role": "assistant", "content": claude_response.content})
                messages.append({
                    "role": "user",
                    "content": "Your previous response was cut off. Please continue exactly from where you left off. Do not include markdown or explanations, just the JSON.",
                })
            elif stop_reason == "refusal":
                raise ValueError("Claude refused to analyze this image due to content policy.")
            else:
                break
                
        response_text = full_response

    try:
        return _parse_json_response(response_text)
    except json.JSONDecodeError as exc:
        raise ValueError(f"OCR returned invalid JSON: {response_text}") from exc