import os
import sys
from typing import Optional, List
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_ollama import ChatOllama


# Pydantic модели
class PersonInfo(BaseModel):
    """Информация о человеке"""
    name: str = Field(description="Имя человека")
    age: Optional[int] = Field(default=None, description="Возраст человека (если указан)")
    profession: str = Field(description="Профессия человека")
    skills: List[str] = Field(default_factory=list, description="Список навыков человека")


class MeetingNotes(BaseModel):
    """Заметки о встрече"""
    date: str = Field(description="Дата встречи")
    participants: List[str] = Field(default_factory=list, description="Список участников встрече")
    topics: List[str] = Field(default_factory=list, description="Обсуждаемые темы")
    decisions: List[str] = Field(default_factory=list, description="Принятые решения")
    next_steps: List[str] = Field(default_factory=list, description="Следующие шаги")


def get_llm():
    """Получить LLM (OpenAI или Ollama)"""
    if os.getenv("OPENAI_API_KEY"):
        return ChatOpenAI(model="gpt-4o-mini", temperature=0)
    return ChatOllama(model="llama3.2", temperature=0)


def create_person_chain():
    """Цепочка для извлечения информации о человеке"""
    parser = PydanticOutputParser(pydantic_object=PersonInfo)
    prompt = PromptTemplate(
        template="""
Извлеки информацию о человеке из текста.
{format_instructions}

Текст: {text}
""",
        input_variables=["text"],
        partial_variables={"format_instructions": parser.get_format_instructions()}
    )
    llm = get_llm()
    return prompt | llm | parser


def create_meeting_chain():
    """Цепочка для извлечения заметок о встрече"""
    parser = PydanticOutputParser(pydantic_object=MeetingNotes)
    prompt = PromptTemplate(
        template="""
Извлеки информацию о встрече из текста.
{format_instructions}

Текст: {text}
""",
        input_variables=["text"],
        partial_variables={"format_instructions": parser.get_format_instructions()}
    )
    llm = get_llm()
    return prompt | llm | parser


def classify_text(text: str) -> str:
    """Определить тип текста: 'person' или 'meeting'"""
    text_lower = text.lower()
    
    # Эвристика для определения типа текста
    meeting_keywords = ["встреча", "meeting", "участники", "participants", "решения", 
                      "decisions", "темы", "topics", "следующие шаги", "next steps",
                      "провели", "прошла", "дата", "date"]
    
    person_keywords = ["год", "лет", "years old", "возраст", "age", "профессия",
                     "profession", "навыки", "skills", "умеет", "знает"]
    
    meeting_score = sum(1 for kw in meeting_keywords if kw in text_lower)
    person_score = sum(1 for kw in person_keywords if kw in text_lower)
    
    if meeting_score >= person_score and meeting_score > 0:
        return "meeting"
    return "person"


def extract_structured_data(text: str) -> BaseModel:
    """Извлечь структурированные данные из текста"""
    text_type = classify_text(text)
    
    if text_type == "meeting":
        chain = create_meeting_chain()
    else:
        chain = create_person_chain()
    
    return chain.invoke({"text": text})


def main():
    """CLI для демонстрации"""
    # Встроенные примеры
    person_example = "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
    meeting_example = "Встреча 28.05.2026. Участники: Иван, Мария, Петр. Темы: план проекта, бюджет. Решения: утвердить бюджет. Следующие шаги: подготовить ТЗ."
    
    print("=== Структурированный вывод (Pydantic) ===\n")
    
    if len(sys.argv) > 1:
        # Ввод с аргументом командной строки
        text = " ".join(sys.argv[1:])
    else:
        # Выбор примера
        print("Выберите пример:")
        print("1. Информация о человеке")
        print("2. Заметки о встрече")
        print("3. Ввести свой текст")
        
        choice = input("Ваш выбор (1-3): ").strip()
        
        if choice == "1":
            text = person_example
        elif choice == "2":
            text = meeting_example
        elif choice == "3":
            text = input("Введите текст: ")
        else:
            print("Неверный выбор, используется пример 1")
            text = person_example
    
    print(f"\nВход: {text}\n")
    
    try:
        result = extract_structured_data(text)
        print("Результат:")
        print(result.model_dump_json(indent=2))
        print(f"\nТип: {type(result).__name__}")
    except Exception as e:
        print(f"Ошибка: {e}")


if __name__ == "__main__":
    main()