import os
import anthropic
import json
import base64

def extract(image_bytes: bytes, prompt: str):
    api_key = os.environ.get("CLAUDE_API_KEY")
    if not api_key:
        return None

    client = anthropic.Anthropic(api_key=api_key)
    base64_image = base64.b64encode(image_bytes).decode('utf-8')
    
    # 2026 Target models
    models_to_try = ["claude-3-5-haiku-latest", "claude-3-5-sonnet-latest"]

    for model_name in models_to_try:
        try:
            print(f"Claude Attempt: {model_name}")
            message = client.messages.create(
                model=model_name,
                max_tokens=1024,
                messages=[{
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": "image/jpeg",
                                "data": base64_image,
                            },
                        },
                        {"type": "text", "text": prompt}
                    ],
                }]
            )
            text = message.content[0].text.strip()
            
            if "```json" in text:
                text = text.split("```json")[1].split("```")[0].strip()
            
            return json.loads(text)
        except Exception as e:
            print(f"Claude {model_name} failed: {e}")
            continue
    return None