#!/usr/bin/env python3

"""
extract-financial-disclosure.py
Parses and extracts structured data from the screenshot at the given URL:
https://gist.github.com/user-attachments/assets/e430e76a-2519-43fa-a370-85a584b816b6

The page comes from page 5 of 24; Schedule III, of the full financial
disclosure report found here:
https://gist.github.com/user-attachments/assets/e430e76a-2519-43fa-a370-85a584b816b6

(the page was manually rotated 90 degrees from its original orientation in the scanned document)

This script assumes your API key is set up in the default way,
  i.e. environment variable: $OPENAI_API_KEY
  https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety
"""
import base64
import json
from openai import OpenAI
from pathlib import Path
from pydantic import BaseModel, Field
from typing import Union, Literal

INPUT_URL = "https://gist.github.com/user-attachments/assets/52c5c8f5-886f-45fe-a338-d1cd3e36ecc8"


# OpenAI examples of Stuctured Output scripts and data definitions
# https://platform.openai.com/docs/guides/structured-outputs/examples?context=ex2


# Define the data structures in Pydantic:
# a Disclosure Report has a list of assets
class Asset(BaseModel):
    owner: Union[Literal['SP', 'DC', 'JT'], None] = Field(description="The leftmost first column of the table")

    asset_name: str = Field(
        description="The name of the asset, the second column of the table"
    )
    asset_value_low: Union[int, None] = Field(
        description="In the third column, the left value of the string and converted to an integer, e.g. '15001' from '$15,001 - $50,000'"
    )
    asset_value_high: Union[int, None] = Field(
        description="In the third column, the right value of the string and converted to an integer, e.g. '50000' from '$15,001 - $50,000'"
    )
    income_type: str = Field(description="The fourth column")

    income_low: Union[int, None] = Field(
        description="In the 5th column, the left value of the string and converted to an integer, e.g. '15001' from '$15,001 - $50,000'.  If the value is enclosed in parentheses, then the income values are meant to be negative"
    )
    income_high: Union[int, None] = Field(
        description="In the 5th column,  the right value of the string and converted to an integer, e.g. '50000' from '$15,001 - $50,000'. If the value is enclosed in parentheses, then the income values are meant to be negative"
    )

    transaction_type: Union[Literal['P', 'S', 'E'], None]

class DisclosureReport(BaseModel):
    assets: list[Asset]

## initialize OpenAI client
client = OpenAI()


# Example of message format for passing in an image via URL
# https://cookbook.openai.com/examples/gpt4o/introduction_to_gpt4o#url-image-processing
input_messages = [
    {"role": "system", "content": "Output the result in JSON format."},
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "Extract the text from this image"},
            {
                "type": "image_url",
                "image_url": {"url": INPUT_URL},
            },
        ],
    },
]

# gpt-4o-mini is cheap and fast and has vision capabilities
response = client.beta.chat.completions.parse(
    response_format=DisclosureReport,
    model="gpt-4o-mini",
    messages=input_messages
)

message = response.choices[0].message

# Print it out in readable format
obj = json.loads(message.content)
print(json.dumps(obj, indent=2))