"""CLI tool for extracting structured data from text using LangChain and Pydantic.

The script supports two schemas:
- PersonInfo
- MeetingNotes

It automatically detects the type of input text and runs the appropriate chain.
"""

import argparse
import sys
from datetime import datetime
from typing import List, Optional

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

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

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

# ---------------------------------------------------------------------------
# Prompt templates
# ---------------------------------------------------------------------------
PERSON_PROMPT = PromptTemplate(
    input_variables=["text"],
    template=(
        """
        Извлечь из следующего текста информацию о человеке в формате JSON, соответствующем схеме PersonInfo.
        Текст: {text}
        JSON: """
    ),
)

MEETING_PROMPT = PromptTemplate(
    input_variables=["text"],
    template=(
        """
        Извлечь из следующего текста информацию о встрече в формате JSON, соответствующем схеме MeetingNotes.
        Текст: {text}
        JSON: """
    ),
)

# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------

def detect_schema(text: str) -> str:
    """Простейшая эвристика: если в тексте встречается слово "встреча" или "meeting" – считаем это встречей.
    Иначе – человек.
    """
    lowered = text.lower()
    if "встреча" in lowered or "meeting" in lowered:
        return "meeting"
    return "person"


def build_chain(schema: str):
    if schema == "person":
        parser = PydanticOutputParser(pydantic_object=PersonInfo)
        chain = PERSON_PROMPT | ChatOpenAI(temperature=0) | parser
    else:
        parser = PydanticOutputParser(pydantic_object=MeetingNotes)
        chain = MEETING_PROMPT | ChatOpenAI(temperature=0) | parser
    return chain

# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Extract structured data from text.")
    parser.add_argument("--text", required=True, help="Input text to parse")
    args = parser.parse_args()

    schema = detect_schema(args.text)
    chain = build_chain(schema)
    result = chain.invoke({"text": args.text})

    # result is a Pydantic model instance
    print(result.model_dump(indent=2))


if __name__ == "__main__":
    main()