main.py

import os
import json
from google import genai
from google.genai import types
from PIL import Image
from dotenv import load_dotenv
from pydantic import BaseModel, Field

load_dotenv(override=True)

# 1. Define a single data point
class DailyStep(BaseModel):
    date: str = Field(description="The date in 'DD/MM' format, or the string 'value' if no date is visible.")
    steps: int = Field(description="The total steps walked as a pure integer.")

# 2. Wrap it in a list container (This avoids the dictionary error entirely)
class StepData(BaseModel):
    steps_log: list[DailyStep] = Field(description="A list containing the extracted step entries from the dashboard.")

def extract_step_count(image_path):
    try:
        img = Image.open(image_path)
    except Exception as e:
        return f"Error opening image: {e}"

    print("Analyzing dashboard with Gemini using Developer-compatible Schema...")
    
    try:
        client = genai.Client()
        
        system_instruction = (
            "You are a precise data extraction AI. Look at fitness dashboard screenshots "
            "and extract step counts. Match the total steps walked with their respective dates. "
            "Never include commas, letters, or symbols in the step value—it must be a pure number."
        )

        response = client.models.generate_content(
            model='gemini-2.5-flash',
            contents=[img, "Extract the step data from this image according to the schema."],
            config=types.GenerateContentConfig(
                system_instruction=system_instruction,
                response_mime_type="application/json",
                response_schema=StepData,
                temperature=0.1,
            ),
        )
        
        # Parse the raw JSON string
        raw_json = response.text
        data_dict = json.loads(raw_json)
        
        # 3. Optional: Convert it back into a flat dictionary format if your code needs it
        # Transforming [{"date": "24/05", "steps": 15185}] -> {"24/05": 15185}
        flat_results = {}
        for entry in data_dict.get("steps_log", []):
            flat_results[entry["date"]] = entry["steps"]
            
        return flat_results
        
    except Exception as e:
        return f"API Error: {e}"

if __name__ == "__main__":
    target_image = "google_fit_screenshot.png" 
    extracted_data = extract_step_count(target_image)
    
    print("\n--- Final Clean Dictionary ---")
    print(extracted_data)