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

# LLM placeholder
llm = ChatOpenAI(
    model="openai/gpt-oss-20b",
    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
    temperature=0.7,
)

class PersonInfo(BaseModel):
    name: str = Field(description="Full name")
    age: int | None = Field(default=None, description="Age in years")
    profession: str = Field(description="Job title")
    skills: list[str] = Field(description="List of technical skills")

class MeetingNotes(BaseModel):
    date: str = Field(description="Meeting date in ISO format")
    participants: list[str] = Field(description="Names of attendees")
    topics: list[str] = Field(description="Discussion topics")
    decisions: list[str] = Field(description="Decisions made")
    next_steps: list[str] = Field(description="Action items")

# Prompt templates
person_prompt = PromptTemplate(
    input_variables=["text"],
    template=(
        "Extract a PersonInfo object from the following text. "
        "Return only JSON matching the schema.\n\n"
        "{format_instructions}\n\nText: {text}"
    ),
)
meeting_prompt = PromptTemplate(
    input_variables=["text"],
    template=(
        "Extract a MeetingNotes object from the following text. "
        "Return only JSON matching the schema.\n\n"
        "{format_instructions}\n\nText: {text}"
    ),
)

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

def route_and_parse(text: str):
    # Simple heuristic: presence of "meeting" or date-like pattern
    if "meeting" in text.lower() or any(c.isdigit() for c in text[:10]):
        prompt = meeting_prompt.partial(
            format_instructions=meeting_parser.get_format_instructions()
        )
        chain = prompt | llm | meeting_parser
        result = chain.invoke({"text": text})
        return MeetingNotes(**result)
    else:
        prompt = person_prompt.partial(
            format_instructions=person_parser.get_format_instructions()
        )
        chain = prompt | llm | person_parser
        result = chain.invoke({"text": text})
        return PersonInfo(**result)

def main():
    examples = [
        "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.",
        "Meeting on 2024-05-27 with Alice and Bob. Topics: budget, timeline. Decisions: approve Q3 plan. Next steps: send email to stakeholders."
    ]
    if len(sys.argv) > 1:
        inputs = [" ".join(sys.argv[1:])]
    else:
        inputs = examples
    for txt in inputs:
        print("\nInput:", txt)
        obj = route_and_parse(txt)
        print(obj.model_dump())

if __name__ == "__main__":
    main()