"""
Full Expression Algebra: All Operators Together

Converted from cookbook example: 34_full_algebra.py

Usage:
    cd examples
    adk web full_algebra
"""

from pydantic import BaseModel
from adk_fluent import Agent, S, Pipeline, until
from adk_fluent._base import _FallbackBuilder
from dotenv import load_dotenv

load_dotenv()  # loads .env from examples/ (copy .env.example -> .env)


class Report(BaseModel):
    title: str
    body: str
    confidence: float


# The complete proof: all operators compose into one expression
#
#   |   parallel research
#   >>  sequence
#   fn  state transform
#   @   typed output
#   //  fallback
#   *   conditional loop
#   S   state transforms
pipeline = (
    # Step 1: Parallel research (|)
    (
        Agent("web").model("gemini-2.5-flash").instruct("Search the web for information.")
        | Agent("scholar").model("gemini-2.5-flash").instruct("Search academic papers.")
    )
    # Step 2: Merge research results (S transform via >>)
    >> S.merge("web", "scholar", into="research")
    # Step 3: Write with typed output (@) and fallback (//)
    >> Agent("writer").model("gemini-2.5-flash").instruct("Write a report.")
    @ Report
    // Agent("writer_b").model("gemini-2.5-pro").instruct("Write a report.")
    @ Report
    # Step 4: Quality loop (* until)
    >> (
        Agent("critic").model("gemini-2.5-flash").instruct("Score the report.").save_as("confidence")
        >> Agent("reviser").model("gemini-2.5-flash").instruct("Improve the report.")
    )
    * until(lambda s: s.get("confidence", 0) >= 0.85, max=4)
)

# Sub-expression reuse — immutable operators make this safe
review = Agent("reviewer").model("gemini-2.5-flash").instruct("Review quality.") >> Agent("scorer").model(
    "gemini-2.5-flash"
).instruct("Score.").save_as("score")
quality_gate = until(lambda s: float(s.get("score", 0)) > 0.8, max=3)

# Same sub-expression in two independent pipelines
pipeline_a = (
    Agent("writer_a").model("gemini-2.5-flash").instruct("Write version A.")
    >> review * quality_gate
    >> S.rename(score="final_score_a")
)

pipeline_b = (
    Agent("writer_b").model("gemini-2.5-flash").instruct("Write version B.")
    >> review * quality_gate
    >> S.rename(score="final_score_b")
)

root_agent = pipeline_b.build()
