import os
import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ---------- Pydantic models ----------
class PersonInfo(BaseModel):
    name: str = Field(description="Full name of the person")
    age: int | None = Field(default=None, description="Age in years, optional")
    profession: str = Field(description="Current profession or job title")
    skills: list[str] = Field(description="List of professional skills")

class MeetingNotes(BaseModel):
    date: str = Field(description="Meeting date in ISO format (YYYY-MM-DD)")
    participants: list[str] = Field(description="Names of participants")
    topics: list[str] = Field(description="Main discussion topics")
    decisions: list[str] = Field(description="Decisions made during the meeting")
    next_steps: list[str] = Field(description="Action items for next steps")

# ---------- Prompt templates ----------
person_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
You are a data extraction assistant. Extract the following information from the given text and output a JSON object that matches the PersonInfo schema.

Text: {text}

Output must be a valid JSON object with fields: name, age, profession, skills.
"""
)

meeting_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
You are a data extraction assistant. Extract the following information from the given text and output a JSON object that matches the MeetingNotes schema.

Text: {text}

Output must be a valid JSON object with fields: date, participants, topics, decisions, next_steps.
"""
)

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

# ---------- Chains ----------
person_chain = person_prompt | llm | person_parser
meeting_chain = meeting_prompt | llm | meeting_parser

# ---------- Backend ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ---------- Agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a structured data extraction agent. Decide whether the input text is about a person or a meeting and return the parsed JSON accordingly.",
)

# ---------- Helper for routing ----------
async def route_and_parse(text: str):
    # Simple heuristic: if the word "meeting" or "встреча" appears, treat as meeting
    if "meeting" in text.lower() or "встреча" in text.lower():
        result = await meeting_chain.ainvoke({"text": text})
        return MeetingNotes(**result)
    else:
        result = await person_chain.ainvoke({"text": text})
        return PersonInfo(**result)

# ---------- CLI ----------
async def main():
    examples = {
        "person": "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.",
        "meeting": "Встреча 2026-06-10. Участники: Иван, Мария. Темы: проект X, бюджет. Решения: утвердить план. Next steps: подготовить презентацию.",
    }
    print("Выберите пример: 1 - человек, 2 - встреча, 3 - ввод с клавиатуры")
    choice = input("> ")
    if choice == "1":
        text = examples["person"]
    elif choice == "2":
        text = examples["meeting"]
    else:
        text = input("Введите текст: ")
    parsed = await route_and_parse(text)
    print("\nРезультат (model_dump):")
    print(parsed.model_dump(indent=2))
    print("\nКраткая сводка:")
    if isinstance(parsed, PersonInfo):
        print(f"{parsed.name}, {parsed.age or 'неизвестно'} лет, {parsed.profession}. Навыки: {', '.join(parsed.skills)}")
    else:
        print(f"Дата: {parsed.date}\nУчастники: {', '.join(parsed.participants)}\nТемы: {', '.join(parsed.topics)}\nРешения: {', '.join(parsed.decisions)}\nСледующие шаги: {', '.join(parsed.next_steps)}")

if __name__ == "__main__":
    asyncio.run(main())