import os
import argparse
from datetime import datetime
from typing import List, Union

from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.runnables import Runnable
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

# ------------------------------------------------------------------
# 1. Pydantic модели
# ------------------------------------------------------------------
class PersonInfo(BaseModel):
    name: str = Field(..., description="Имя человека")
    age: int | None = Field(None, description="Возраст (необязательно)")
    profession: str = Field(..., description="Профессия")
    skills: List[str] = Field(..., description="Список навыков")

class MeetingNotes(BaseModel):
    date: datetime = Field(..., description="Дата встречи в формате YYYY-MM-DD")
    participants: List[str] = Field(..., description="Список участников")
    topics: List[str] = Field(..., description="Темы обсуждения")
    decisions: List[str] = Field(..., description="Принятые решения")
    next_steps: List[str] = Field(..., description="Следующие шаги")

# ------------------------------------------------------------------
# 2. Шаблоны и парсеры
# ------------------------------------------------------------------
person_prompt = PromptTemplate(
    input_variables=["text"],
    template=(
        "Найди в тексте информацию о человеке и выведи её в формате JSON, "
        "соответствующем схеме PersonInfo.\n"
        "Текст: {text}\n"
        "Формат: {format_instructions}"
    ),
)

meeting_prompt = PromptTemplate(
    input_variables=["text"],
    template=(
        "Найди в тексте информацию о встрече и выведи её в формате JSON, "
        "соответствующем схеме MeetingNotes.\n"
        "Текст: {text}\n"
        "Формат: {format_instructions}"
    ),
)

person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)

person_chain = (
    person_prompt | ChatOpenAI(temperature=0) | person_parser
)
meeting_chain = (
    meeting_prompt | ChatOpenAI(temperature=0) | meeting_parser
)

# ------------------------------------------------------------------
# 3. Выбор схемы
# ------------------------------------------------------------------
def choose_chain(text: str) -> Runnable:
    """Определяем, какой тип данных в тексте."""
    lowered = text.lower()
    # простая эвристика: наличие ключевых слов
    if any(word in lowered for word in ("встреча", "meeting", "дата", "participants", "topics")):
        return meeting_chain
    return person_chain

# ------------------------------------------------------------------
# 4. CLI
# ------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Structured output extractor")
    parser.add_argument(
        "--text",
        type=str,
        help="Текст для обработки. Если не задан, используется пример.",
    )
    args = parser.parse_args()

    if args.text:
        input_text = args.text
    else:
        # пример для PersonInfo
        input_text = (
            "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
        )
        # пример для MeetingNotes
        # input_text = (
        #     "Дата: 2024-05-01. Участники: Анна, Боб. Темы: проект X, бюджет. "
        #     "Решения: утвердить бюджет. Next steps: подготовить презентацию."
        # )

    chain = choose_chain(input_text)
    result = chain.invoke({"text": input_text})

    # result уже валидированный объект Pydantic
    print("\nВыполнено. Валидированный объект:")
    print(result.model_dump(indent=2))

    # Краткая сводка
    if isinstance(result, PersonInfo):
        summary = (
            f"Человек: {result.name}, "
            f"возраст: {result.age if result.age is not None else 'не указан'}, "
            f"профессия: {result.profession}, "
            f"навыки: {', '.join(result.skills)}"
        )
    else:
        summary = (
            f"Встреча: {result.date.date()}, "
            f"участники: {', '.join(result.participants)}, "
            f"темы: {', '.join(result.topics)}, "
            f"решения: {', '.join(result.decisions)}, "
            f"следующие шаги: {', '.join(result.next_steps)}"
        )
    print("\nСводка:")
    print(summary)

if __name__ == "__main__":
    # Для работы с OpenAI требуется переменная окружения OPENAI_API_KEY
    if "OPENAI_API_KEY" not in os.environ:
        raise RuntimeError("Необходимо задать переменную окружения OPENAI_API_KEY")
    main()