"""
lambdagent Hello World

从最简单开始，逐步展示 11 个 Lambda 构造。
每一步都标注 Lambda 演算对应。

用法:
    export ANTHROPIC_API_KEY=sk-...
    python hello.py
"""

import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from lambdagent import (
    Lam, Tool, Compose, If, Loop, Pair, Fst, Snd,
    Route, Guard, Context, Dataset,
)
from lambdagent.extensions import Memory


def main():
    ctx = Context()

    # ════════════════════════════════════════════
    # 1. Lam — λ 抽象
    #    λx. LLM(x)
    # ════════════════════════════════════════════
    print("─── 1. Lam (λ abstraction) ───")

    hello = Lam(
        "hello",
        "You are a friendly assistant. Reply in one short sentence.",
        max_tokens=64,
    )
    print(f"  λ> {hello('Hello!', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 2. Tool — 原语 / Oracle
    #    λx. f(x)  where f is a Python function
    # ════════════════════════════════════════════
    print("─── 2. Tool (primitive) ───")

    shout = Tool("shout", lambda x: str(x).upper() + "!!!")
    print(f"  λ> {shout('hello world', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 3. Compose — 函数组合
    #    f >> g = λx. g(f(x))
    # ════════════════════════════════════════════
    print("─── 3. Compose (f >> g) ───")

    pipeline = hello >> shout   # 先让 LLM 回复，再全大写
    print(f"  λ> {pipeline('Hi there', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 4. Dataset → Lam — 数据集变函数
    #    D.to_lam() = λx. LLM_{prompt=encode(D)}(x)
    # ════════════════════════════════════════════
    print("─── 4. Dataset → Lam ───")

    succ = Dataset(
        examples=[("0","1"), ("1","2"), ("2","3"), ("9","10")],
        description="Given n, output n+1. Output ONLY the number.",
    ).to_lam("SUCC")

    print(f"  SUCC(41) = {succ('41', ctx)}")
    print(f"  SUCC(99) = {succ('99', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 5. Pair / Fst / Snd — Church 对
    #    PAIR = λa.λb.λf. f a b
    # ════════════════════════════════════════════
    print("─── 5. Pair / Fst / Snd ───")

    both = Pair(hello, shout)
    result = both("Hey", ctx)
    print(f"  Pair result: {result}")
    print(f"  Fst: {Fst()('(a, b) tuple test: ' + str(result), ctx) if False else result[0]}")
    print(f"  Snd: {result[1]}")
    print()

    # ════════════════════════════════════════════
    # 6. If — Church 条件
    #    IF c t e = c t e
    # ════════════════════════════════════════════
    print("─── 6. If (Church boolean) ───")

    branch = If(
        cond=lambda x: int(x) > 10,
        then_=Tool("big", lambda x: f"{x} is BIG"),
        else_=Tool("small", lambda x: f"{x} is small"),
    )
    print(f"  If(5 > 10):  {branch('5', ctx)}")
    print(f"  If(42 > 10): {branch('42', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 7. Loop — Y 组合子
    #    Y(λself.λx. if done then x else self(body(x)))
    # ════════════════════════════════════════════
    print("─── 7. Loop (Y combinator) ───")

    counter = Tool("inc", lambda x: str(int(x) + 1))
    loop = Loop(
        body=counter,
        condition=lambda result, step: int(result) >= 5,
        max_steps=20,
    )
    print(f"  Y(inc)(0) until ≥5: {loop('0', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 8. Route — 广义 Church 布尔 (CASE)
    #    CASE classifier(x) [(l₁,a₁), (l₂,a₂), ...]
    # ════════════════════════════════════════════
    print("─── 8. Route (CASE) ───")

    classifier = Lam(
        "classify",
        "Classify the input as exactly one word: 'math' or 'greeting'. Output ONLY that word.",
        max_tokens=8,
    )
    router = Route(
        classifier=classifier,
        routes={
            "math": Tool("math_handler", lambda x: f"[MATH] {x}"),
            "greeting": Tool("greet_handler", lambda x: f"[GREET] Hello! {x}"),
        },
        default=Tool("default", lambda x: f"[DEFAULT] {x}"),
    )
    print(f"  Route('2+2'):  {router('2+2', ctx)}")
    print(f"  Route('Hi!'):  {router('Hi!', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 9. Guard — 依赖类型 {x:T | P(x)}
    #    if P(result) then result else retry
    # ════════════════════════════════════════════
    print("─── 9. Guard (dependent type) ───")

    strict_num = Guard(
        agent=succ,
        validator=lambda x: str(x).strip().isdigit(),
        retry=2,
    )
    print(f"  Guard(SUCC(7)): {strict_num('7', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 10. Memory — 环境扩展 Γ' = Γ ∪ store
    # ════════════════════════════════════════════
    print("─── 10. Memory (Γ extension) ───")

    greeter = Lam(
        "greeter",
        "Greet the user. If you know their name from memory, use it. One sentence.",
        max_tokens=64,
    )
    stateful = Memory(greeter, store={"user_name": "Alice", "mood": "happy"})
    print(f"  Memory(greeter): {stateful('Hello!', ctx)}")
    print()

    # ════════════════════════════════════════════
    # 11. 组合: 完整 Agent pipeline
    #     extract >> Pair(analyze, critique) >> synthesize >> Loop(refine)
    # ════════════════════════════════════════════
    print("─── 11. Full pipeline ───")

    extract = Lam("extract", "Extract the key claim from this text. One sentence.", max_tokens=64)
    analyze = Lam("analyze", "Is this claim true or false? One sentence.", max_tokens=64)
    combine = Tool("combine", lambda pair: f"Claim: {pair[0]}\nAnalysis: {pair[1]}")

    full = extract >> Pair(analyze, shout) >> combine
    result = full("The Earth orbits the Sun in approximately 365 days.", ctx)
    print(f"  {result}")
    print()

    # ════════════════════════════════════════════
    # β-规约追踪
    # ════════════════════════════════════════════
    print("─── β-reduction trace ───")
    ctx.print_trace()
    print(f"\nTotal: {len(ctx.trace)} β-reductions")


if __name__ == "__main__":
    main()