"""runnable_lambda_demo.py — custom Python in a chain
Run with venv active:  python runnable_lambda_demo.py
"""
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableLambda

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini", temperature=0)
parser = StrOutputParser()

pick_prompt = ChatPromptTemplate.from_template(
    "Name one HTML tag used for {feature}. Reply with only the tag name, no brackets."
)
pick_chain = pick_prompt | model | parser


def normalize_tag(raw: str) -> dict:
    """Strip whitespace and angle brackets from the model's tag string."""
    clean = raw.strip().strip("<>").lower()
    return {"tag": clean}


explain_prompt = ChatPromptTemplate.from_template(
    "What does the HTML <{tag}> tag do? One short sentence."
)

# RunnableLambda runs your function between chain steps
chain = pick_chain | RunnableLambda(normalize_tag) | explain_prompt | model | parser

features = ["links", "images"]

for feature in features:
    explanation = chain.invoke({"feature": feature})
    print(f"Feature: {feature}")
    print(f"Explanation: {explanation}\n")