import os
import sys
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from models import PersonInfo, MeetingNotes

# Choose LLM
llm = ChatOpenAI(temperature=0)

# Prepare parser for each model
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)

# Prompt templates
person_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
    Extract structured data about a person from the following text.
    The output must follow the JSON schema:
    {schema}
    Text: {text}
    """,
    partial_variables={"schema": person_parser.get_format_instructions()},
)

meeting_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
    Extract structured data about a meeting from the following text.
    The output must follow the JSON schema:
    {schema}
    Text: {text}
    """,
    partial_variables={"schema": meeting_parser.get_format_instructions()},
)

# Helper to detect type

def detect_type(text: str) -> str:
    # Simple heuristic: if contains "meeting" or "participants" -> meeting, else person
    lower = text.lower()
    if "meeting" in lower or "participants" in lower or "topics" in lower:
        return "meeting"
    return "person"

def main():
    if len(sys.argv) > 1:
        input_text = " ".join(sys.argv[1:])
    else:
        print("Enter text (or press Ctrl-D to exit):")
        input_text = sys.stdin.read().strip()
        if not input_text:
            print("No input provided.")
            return

    t = detect_type(input_text)
    if t == "person":
        chain = person_prompt | llm | person_parser
        result = chain.invoke({"text": input_text})
        print("\nParsed PersonInfo:\n", result.model_dump(indent=2))
    else:
        chain = meeting_prompt | llm | meeting_parser
        result = chain.invoke({"text": input_text})
        print("\nParsed MeetingNotes:\n", result.model_dump(indent=2))

if __name__ == "__main__":
    main()