import asyncio
import io
import json
import sys

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

import os

from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI

from agent.lead_stage_agent.prompts import LEAD_STAGE_CLASSIFIER_PROMPT

load_dotenv()

async def main():
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
        temperature=0.2
    ).with_structured_output(
        schema={
            "type": "object",
            "properties": {
                "inferred": {
                    "type": "object",
                    "properties": {
                        "product_line_vn": {"type": "array", "items": {"type": "string"}},
                        "gender_by_product": {"type": "string"},
                        "keywords": {"type": "array", "items": {"type": "string"}},
                    }
                },
                "literal": {
                    "type": "object",
                    "properties": {
                        "raw_text": {"type": "string"}
                    }
                }
            }
        }
    )
    
    prompt = LEAD_STAGE_CLASSIFIER_PROMPT + "\n\nChat history:\nUser: tìm sịp nam cực sinh"
    
    result = await llm.ainvoke(prompt)
    print(json.dumps(result, indent=2, ensure_ascii=False))

if __name__ == '__main__':
    asyncio.run(main())
