# app/agents/receipt_agent.py
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import PydanticOutputParser
from langchain.prompts import PromptTemplate
from ..core.llm import llm
from ..schemas import ReceiptInfo
from typing import List

parser = PydanticOutputParser(pydantic_object=ReceiptInfo)

receipt_template = """
You are a highly accurate receipt information extraction robot.
The document's primary language is '{language}'.
Please extract all key information from the following receipt images, which belong to a single document.
If some information is not present in the images, leave it as null.
Please strictly follow the JSON format below, without adding any extra explanations or comments.

{format_instructions}
"""

receipt_prompt = PromptTemplate(
    template=receipt_template,
    input_variables=["language"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)


async def agent_extract_receipt_info(images_base64: List[str], language: str) -> ReceiptInfo:
    """Agent 2: Extracts receipt information from a list of images, aware of the document's language."""
    print(f"--- [Agent 2] Calling multimodal LLM to extract receipt info (Language: {language})...")

    prompt_text = await receipt_prompt.aformat(language=language)

    content_parts = [{"type": "text", "text": prompt_text}]
    for image_base64 in images_base64:
        content_parts.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{image_base64}"},
        })

    msg = HumanMessage(content=content_parts)

    chain = llm | parser
    receipt_info = await chain.ainvoke([msg])
    return receipt_info