"""
agentbudget — OpenAI Usage Example

This example demonstrates how to track the cost of OpenAI API calls
using agentbudget's CostSession context manager.

Run:
    pip install agentbudget[openai]
    export OPENAI_API_KEY=your-key
    python examples/openai_example.py
"""

from agentbudget import CostSession

# NOTE: Replace with your actual OpenAI client
# from openai import OpenAI
# client = OpenAI()


def run_pipeline() -> None:
    """Demonstrate cost tracking for a multi-step OpenAI pipeline."""

    # Simulated responses (replace with real OpenAI calls)
    class FakeUsage:
        def __init__(self, prompt_tokens: int, completion_tokens: int) -> None:
            self.prompt_tokens = prompt_tokens
            self.completion_tokens = completion_tokens

    class FakeResponse:
        def __init__(self, model: str, prompt_tokens: int, completion_tokens: int) -> None:
            self.model = model
            self.usage = FakeUsage(prompt_tokens, completion_tokens)

    with CostSession(name="content-ideas-pipeline", budget=0.05) as session:
        # Step 1: Generate ideas
        # response = client.chat.completions.create(
        #     model="gpt-4o",
        #     messages=[{"role": "user", "content": "Give me 5 blog ideas about AI"}]
        # )
        response = FakeResponse("gpt-4o", prompt_tokens=800, completion_tokens=443)
        session.track(response, label="generate")

        # Step 2: Score ideas
        # response = client.chat.completions.create(
        #     model="gpt-4o",
        #     messages=[{"role": "user", "content": "Score each idea 1-10"}]
        # )
        response = FakeResponse("gpt-4o", prompt_tokens=1200, completion_tokens=908)
        session.track(response, label="score")

        # Step 3: Rewrite best idea
        # response = client.chat.completions.create(
        #     model="gpt-4o-mini",
        #     messages=[{"role": "user", "content": "Rewrite the best idea"}]
        # )
        response = FakeResponse("gpt-4o-mini", prompt_tokens=500, completion_tokens=391)
        session.track(response, label="rewrite")

    # You can also get the summary as a string for logging
    # summary_str = session.summary()


if __name__ == "__main__":
    run_pipeline()
