# app/agents/invoice_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 InvoiceInfo
from typing import List

parser = PydanticOutputParser(pydantic_object=InvoiceInfo)

invoice_template = """
You are an expert data entry clerk AI. Your primary goal is to extract information from an invoice image with the highest possible accuracy.
The document's primary language is '{language}'.

## Instructions:
Carefully analyze the invoice image and extract the following fields according to these specific rules. Do not invent information. If a field is not found or is unclear, follow the specific instruction for that field.

- `invoice_date`: The invoice date. Extract in YYYY-MM-DD format. If unclear, leave as an empty string.
- `invoice_due_date`: The invoice due date.Extract in YYYY-MM-DD format. If unclear, leave as an empty string.
- `invoice_number`: If not found or unclear, leave as an empty string.
- `ocr_number`: The OCR number from the invoice. If not found or unclear, leave as an empty string.
- `supplier_number`: This is the organisation number. If not found or unclear, leave as an empty string.
- `biller_name`: This is the sender's name. If not found or unclear, leave as an empty string.
- `amount`: Extract the final total amount and format it to a decimal number. If not present, leave as null.
- `tax_exclusive_amount`: Extract the the amount excluding taxes and format it to a decimal number. If not present, leave as null.
- `customer_name`: This is the receiver's name. Ensure it is a name and clear any special characters. If not found or unclear, leave as an empty string.
- `customer_address`: This is the receiver's full address. Put it in one line. If not found or unclear, leave as an empty string.
- `customer_address_line`: This is only the street address line from the receiver's address. If not found or unclear, leave as an empty string.
- `customer_address_city`: This is the receiver's city. If not found, try to find any city in the document. If unclear, leave as an empty string.
- `customer_address_country`: This is the receiver's country. If not found, find the country of the extracted city. If unclear, leave as an empty string.
- `customer_address_postal_code`: This is the receiver's postal code. If not found or unclear, leave as an empty string.
- `customer_address_apartment`: This is the receiver's apartment or suite number. If not found or unclear, leave as an empty string.
- `customer_address_region`: This is the receiver's region. If not found, find the region of the extracted city or country. If unclear, leave as an empty string.
- `customer_address_care_of`: This is the receiver's 'care of' (c/o) line. If not found or unclear, leave as an empty string.
- `billo_id`: To find this, think step-by-step: 1. Find the customer_address. 2. Scan the address for a pattern of three letters, an optional space, three digits, an optional dash, and one alphanumeric character (e.g., 'ABC 123-X' or 'DEF 456Z'). 3. If found, extract it. If not found or unclear, leave as an empty string.
- `bank_giro`: If found, extract the bank giro number. It often follows patterns like 'ddd-dddd', 'dddd-dddd', or 'dddddddd #41#'. If not found or unclear, leave as an empty string.
- `plus_giro`: If found, extract the plus giro number. It often follows patterns like 'ddddddd-d #16#', 'ddddddd-d', or 'ddd dd dd-d'. If not found or unclear, leave as an empty string.
- `customer_ssn`: If found, extract the customer social security number (personnummer). It follows the pattern 'YYYYMMDD-XXXX' or 'YYMMDD-XXXX'. If not found or unclear, leave as an empty string.
- `line_items`: Extract all line items from the invoice. For each item, extract the `description`, `quantity`, `unit_price`, and `total_price`. A list of all line items from the invoice. Make sure all of them are extracted. If a value is not present, leave it as null.

## Example:
If the invoice shows a line item "Consulting Services | 2 hours | $100.00/hr | $200.00", the output for that line item should be:
```json
{{
  "description": "Consulting Services",
  "quantity": 2,
  "unit_price": 100.00,
  "total_price": 200.00
}}
```
Your Task:
Now, analyze the provided image and output the full JSON object according to the format below.
{format_instructions}
"""
invoice_prompt = PromptTemplate(
    template=invoice_template,
    input_variables=["language"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

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

    prompt_text = await invoice_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
    invoice_info = await chain.ainvoke([msg])
    return invoice_info