import os
import json
import asyncio
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend
from langchain_core.messages import HumanMessage

# Загрузим переменные окружения (ключ OpenRouter)
load_dotenv()

# 1. Модель карточки задания
class TaskCard(BaseModel):
    title: str = Field(description="Title of the assignment.")
    subject: str = Field(description="Subject or topic.")
    deadline_hint: str = Field(description="Short hint about the deadline.")
    deliverable_type: str = Field(description="What to submit (report, code, presentation, etc.).")
    grading_hints: list[str] = Field(description="List of hints about grading criteria.")

# 2. Парсер
parser = PydanticOutputParser(pydantic_object=TaskCard)

# 3. Шаблон промпта
prompt = PromptTemplate(
    template="""Extract the following fields from the task description:
title: the title of the assignment
subject: the subject or topic
deadline_hint: a short hint about deadline (e.g., 'by Friday')
deliverable_type: what to submit (report, code, presentation, etc.)
grading_hints: list of hints about grading criteria

Return the data in JSON format as follows:
{format_instructions}
The description: {description}
""",
    input_variables=["description"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# 4. LLM (OpenRouter)
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,
)

# 5. Цепочка: промпт → LLM → парсер
chain = prompt | llm | parser

# 6. Инструмент для парсинга
@tool
def parse_task(description: str) -> TaskCard:
    """Parse a task description into structured fields."""
    return chain.invoke({"description": description})

# 7. Бэкенд и агент
backend = FilesystemBackend()
agent = create_deep_agent(
    model=llm,
    tools=[parse_task],
    backend=backend,
    system_prompt="You are a task parsing agent. Use the parse_task tool to extract structured data from the assignment description.",
)

# 8. Основная функция
async def main():
    description = (
        "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. "
        "Оценка: за полноту и за пример кода."
    )

    # Вызов агента
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=description)]},
        {"configurable": {"thread_id": "session-1"}},
    )

    # Последнее сообщение содержит вывод инструмента
    output_text = result["messages"][-1].content

    # Попытка распарсить JSON
    try:
        parsed_dict = json.loads(output_text)
        parsed = TaskCard(**parsed_dict)
    except json.JSONDecodeError:
        # Если вывод не в JSON, используем eval (не рекомендуется в продакшене)
        parsed = eval(output_text)

    # Вывод валидированного объекта
    print("Валидированный объект:")
    print(parsed.model_dump())
    print("\nКраткая сводка:")
    print(f"Задание: {parsed.title}")
    print(f"Тема: {parsed.subject}")
    print(f"Срок: {parsed.deadline_hint}")
    print(f"Сдача: {parsed.deliverable_type}")
    print(f"Критерии оценки: {', '.join(parsed.grading_hints)}")

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