import asyncio
import os
from typing import Any

import httpx
from pydantic import BaseModel
from timbal import Agent, Tool, Workflow
from timbal.errors import bail
from timbal.state import get_run_context
from timbal.tools import WebSearch
from timbal.types import File

# os.environ.pop("ANTHROPIC_API_KEY", None)
os.environ["TIMBAL_API_HOST"] = "api.timbal.ai"
os.environ["TIMBAL_ORG_ID"] = "1"
os.environ["TIMBAL_APP_ID"] = "177"
# os.environ["TIMBAL_PROJECT_ID"] = "60"
os.environ["TIMBAL_LOG_EVENTS"] = "START,OUTPUT,DELTA"


async def slow_tool(query: str):
    await asyncio.sleep(10)


internal_search = Tool(
    name="internal_search",
    description="Searches the internal database for information",
    handler=slow_tool,
)


class Intent(BaseModel):
    intent: str
    confidence: float
    explanation: str


agent = Agent(
    name="test-proxies",
    model="anthropic/claude-haiku-4-5",  # type: ignore
    # model="openai/gpt-5.2",  # type: ignore
    # model="google/gemini-2.5-flash-lite",  # type: ignore
    # tools=[WebSearch(), internal_search],
    # tools=[internal_search],
    model_params={"max_tokens": 128},  # type: ignore
    # output_model=Intent,
)


async def call_remote_agent(prompt: str, parent_id: str | None = None) -> dict[str, Any]:
    async with httpx.AsyncClient(timeout=httpx.Timeout(10.0, read=180.0)) as client:
        res = await client.post(
            f"https://{os.environ['TIMBAL_API_HOST']}/orgs/1/apps/177/runs/collect",
            headers={
                "Authorization": f"Bearer {os.environ['TIMBAL_API_KEY']}",
                "Content-Type": "application/json",
            },
            json={
                # "version_id": "1139",
                "parent_id": parent_id,
                "input": {"prompt": prompt},
            },
        )
        res.raise_for_status()
        return res.json()


def step_a():
    return {"valor": 5}


def step_b(x):
    return {"valor": x * 2}


def step_c(x):
    return {"valor": x * 3}


def step_d(x):
    return {"valor": x}


workflow = (
    Workflow(name="test")
    .step(step_a)
    .step(
        step_b,
        x=lambda: get_run_context().step_span("step_a").output.get("valor"),
        when=lambda: get_run_context().step_span("step_a").output.get("valor") > 10,
    )
    .step(
        step_c,
        x=lambda: get_run_context().step_span("step_a").output.get("valor"),
        when=lambda: get_run_context().step_span("step_a").output.get("valor") > 10,
    )
    .step(
        step_d,
        x=lambda: get_run_context().step_span("step_b").output.get("valor")
        + get_run_context().step_span("step_c").output.get("valor"),
    )
)


async def main():
    # await agent(prompt="when does fcbarcelona play next?", max_tokens=1024).collect()
    # await agent(prompt="hola", max_tokens=1024).collect()
    # await agent(prompt=[File.validate("~/Downloads/Transfer_Confirmation_EUR-USD_02-ene-2026_12.08.55.pdf")]).collect()
    # await agent(prompt="how much were the fees?").collect()
    # await agent(prompt=[File.validate("~/Downloads/Timbal AI - Plan de Incentivos + Anexos vfinal-1.pdf")]).collect()
    # parent_id = None
    # while True:
    #     prompt = input("User: ")
    #     if prompt.strip() == "q":
    #         break
    #     await agent(prompt=prompt).collect()
    # res_json = await call_remote_agent(prompt=prompt, parent_id=parent_id)
    # print(res_json)
    # parent_id = res_json["run_id"]
    await workflow().collect()


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