"""output_formatting_demo.py — plain string vs structured output
Run with venv active:  python output_formatting_demo.py
"""
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
from pydantic import BaseModel, Field

load_dotenv()

class HtmlTagInfo(BaseModel):
    tag: str = Field(description="Tag name without angle brackets, e.g. a")
    purpose: str = Field(description="What the tag does in one short sentence")
    example: str = Field(description="A minimal HTML example using the tag")

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

tags = ["a", "img"]

# Plain string output (quotes around tag name keep "img" unambiguous)
str_chain = (
    ChatPromptTemplate.from_template(
        "What does the HTML '{tag}' element do? One short sentence."
    )
    | model
    | StrOutputParser()
)

print("=== StrOutputParser ===")
for tag in tags:
    text = str_chain.invoke({"tag": tag})
    print(f"Tag: <{tag}>")
    # langchain-core 1.0+ may return TextAccessor — a str subclass you use like a string
    print(f"Type:  {type(text).__name__}")
    print(f"is str: {isinstance(text, str)}")
    print(f"Value: {text}")
    print()

# Structured object output
parser = PydanticOutputParser(pydantic_object=HtmlTagInfo)
struct_chain = (
    ChatPromptTemplate.from_template(
        "Describe the HTML '{tag}' element.\n{format_instructions}"
    ).partial(format_instructions=parser.get_format_instructions())
    | model
    | parser
)

print("=== PydanticOutputParser ===")
for tag in tags:
    info = struct_chain.invoke({"tag": tag})
    print(f"Tag: <{tag}>")
    print(f"Type:  {type(info).__name__}")
    print(f"  tag:     {info.tag}")
    print(f"  purpose: {info.purpose}")
    print(f"  example: {info.example}")
    print()