"""
Get a company's PitchBook ID

Designed to run inside a Zapier "Run Python" step (or any environment that
exposes an `input_data` dict). Given a company name and/or website pulled
from Attio, query PitchBook's search API and return the matching record so
downstream steps can grab the PitchBook ID.
"""

import requests

api_key = input_data['api_key']                 # PB API key
company_name = input_data.get('company_name')   # from Attio
company_website = input_data.get('company_website')  # from Attio

# Choose website if available, else name
search_query = company_website if company_website else company_name

headers = {
    'Authorization': f"PB-Token {api_key}"
}

url = f"https://api.pitchbook.com/search?query={search_query}"

resp = requests.get(url, headers=headers)

# If the API itself errors, keep raising (likely auth/network issue)
resp.raise_for_status()

output = resp.json()