from stagehand import Stagehand
import asyncio
from dotenv import load_dotenv
from pydantic import BaseModel, HttpUrl
from typing import Optional
import json

load_dotenv()

class Agency(BaseModel):
    name: str
    url: HttpUrl
    description: Optional[str] = None
    location: Optional[str] = None

class Agencies(BaseModel):
    agencies: list[Agency]

async def main():
    stagehand = Stagehand(
        env="LOCAL",
        headless=False,
        model_name="gpt-4.1-mini"
    )
    all_agencies = []
    await stagehand.init()
    page = stagehand.page
    await page.goto("<url I want to scrape>")
    should_continue = True
    while should_continue:
        agencies = await page.extract("Extract all agencies from the page", schema=Agencies)
        all_agencies.extend([agency.model_dump(mode='json') for agency in agencies.agencies])
        next_page_button = await page.observe("Find the next page button")
        should_continue = has_next_page(next_page_button)
        if should_continue:
            await page.click(next_page_button[0].selector)
    
    print(all_agencies)
    with open("agencies.json", "w") as f:
        json.dump(all_agencies, f)


if __name__ == "__main__":
    asyncio.run(main())