import argparse
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

# 1. 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="Следующие шаги")

# 2. Prompt and parser
prompt_template = """
Extract structured data from the following text.
Use the appropriate schema:
- PersonInfo for a person description
- MeetingNotes for meeting notes
Return only JSON matching the chosen schema.
Text: {text}
"""
prompt = PromptTemplate.from_template(prompt_template)
parser_person = PydanticOutputParser(pydantic_object=PersonInfo)
parser_meeting = PydanticOutputParser(pydantic_object=MeetingNotes)

# 3. Simple type guesser (very naive)
def choose_parser(text: str):
    if "профессия" in text.lower() or "навыки" in text.lower():
        return parser_person, PersonInfo
    else:
        return parser_meeting, MeetingNotes

# 4. Main chain
llm = ChatOpenAI(temperature=0)

def extract(text: str):
    parser, model_cls = choose_parser(text)
    chain = prompt | llm | parser
    result = chain.invoke({"text": text})
    return result, model_cls

# 5. CLI
if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Extract structured data from raw text")
    ap.add_argument("--text", required=True, help="Raw input text")
    args = ap.parse_args()
    out, model_cls = extract(args.text)
    print(out.model_dump(indent=2))