from agentui.pipeline.templating import (
    render_template_simple,
    _best_text_from_outputs,
)

def run_checks():
    # Common context for tests
    ctx = {
        "model": "gpt-x",
        "params": {"max_tokens": 128},
        "chat": {"last_user": "Привет"},
        "OUT": {},
    }
    out_map = {}

    # 1) [[VAR:...]]
    s1 = render_template_simple("Hello [[VAR:chat.last_user]]", ctx, out_map)
    assert s1 == "Hello Привет"

    # 2) {{ ... |default(...) }} when missing
    s2 = render_template_simple("T={{ params.temperature|default(0.7) }}", ctx, out_map)
    assert s2 == "T=0.7"
    # present
    ctx2 = {**ctx, "params": {**ctx["params"], "temperature": 0.4}}
    s3 = render_template_simple("T={{ params.temperature|default(0.7) }}", ctx2, out_map)
    assert s3 == "T=0.4"

    # 3) [[OUT:n1...]] exact path
    out_map = {
        "n1": {
            "result": {
                "choices": [
                    {"message": {"content": "Hi from OpenAI"}}
                ]
            }
        }
    }
    s4 = render_template_simple("[[OUT:n1.result.choices.0.message.content]]", ctx, out_map)
    assert s4 == "Hi from OpenAI"

    # 4) [[OUT1]] short form -> best-effort text
    s5 = render_template_simple("[[OUT1]]", ctx, out_map)
    assert s5 == "Hi from OpenAI"

    # 5) [[PROMPT]] raw fragment insertion
    ctx_prompt = {**ctx, "PROMPT": '"messages": [{"role":"user","content":"Hi"}]'}
    s6 = render_template_simple('{"model":"{{ model }}", [[PROMPT]] }', ctx_prompt, out_map)
    # ensure prompt fragment is inserted without extra quotes
    assert '"messages": [{"role":"user","content":"Hi"}]' in s6
    assert '"model":"gpt-x"' in s6

    # 6) _best_text_from_outputs (Gemini)
    gemini_out = {
        "result": {
            "candidates": [
                {"content": {"role": "model", "parts": [{"text": "Gemini says"}]}}
            ]
        }
    }
    assert _best_text_from_outputs(gemini_out) == "Gemini says"

    # 7) _best_text_from_outputs (Claude)
    claude_out = {
        "result": {
            "content": [
                {"type": "text", "text": "Claude part 1"},
                {"type": "text", "text": "Claude part 2"},
            ]
        }
    }
    assert _best_text_from_outputs(claude_out) == "Claude part 1\nClaude part 2"

    # 8) _best_text_from_outputs (direct response_text)
    direct = {"response_text": "Direct text"}
    assert _best_text_from_outputs(direct) == "Direct text"

    # 9) Mixed braces with OUT
    out_map = {
        "n2": {"result": {"obj": {"value": 42}}},
    }
    s7 = render_template_simple("V={{ OUT.n2.result.obj.value|default(0) }}", ctx, out_map)
    assert s7 == "V=42"


if __name__ == "__main__":
    run_checks()
    print("Templating tests: OK")