import anthropic
import base64
import sys
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

MEDIA_TYPES = {
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".gif": "image/gif",
    ".webp": "image/webp",
}


def get_media_type(image_path: str) -> str:
    ext = Path(image_path).suffix.lower()
    media_type = MEDIA_TYPES.get(ext)
    if not media_type:
        raise ValueError(f"Unsupported image format: {ext}. Use jpg, png, gif, or webp.")
    return media_type


def load_image_b64(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.standard_b64encode(f.read()).decode("utf-8")


def query_claude(image_path: str, prompt: str) -> str:
    client = anthropic.Anthropic()
    media_type = get_media_type(image_path)
    image_data = load_image_b64(image_path)

    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=4096,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_type,
                            "data": image_data,
                        },
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ],
    )

    return next(block.text for block in response.content if block.type == "text")


def main():
    if len(sys.argv) < 3:
        print("Usage: python image_query.py <image_path> \"<prompt>\" [output_file]")
        sys.exit(1)

    image_path = sys.argv[1]
    prompt = sys.argv[2]
    output_file = sys.argv[3] if len(sys.argv) > 3 else "response.txt"

    if not Path(image_path).exists():
        print(f"Error: image file '{image_path}' not found.")
        sys.exit(1)

    print(f"Querying Claude about '{image_path}'...")
    response_text = query_claude(image_path, prompt)

    print("\n--- Response ---")
    print(response_text)

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(response_text)

    print(f"\nSaved to {output_file}")


if __name__ == "__main__":
    main()