import os
import asyncio
from typing import List

from pydantic import BaseModel, Field

from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain_openai import ChatOpenAI
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain.tools import tool
from langchain_core.messages import HumanMessage

# ------------------------------
# Pydantic model for the task card
# ------------------------------
class TaskCard(BaseModel):
    title: str = Field(description="Краткое название задания")
    subject: str = Field(description="Тема или предмет задания")
    deadline_hint: str = Field(description="Подсказка о сроке выполнения, например 'к пятнице' или 'до 12.09'")
    deliverable_type: str = Field(description="Что нужно сдать: отчёт, код, презентация и т.п.")
    grading_hints: List[str] = Field(description="Список критериев оценки, упомянутых в тексте")

# ------------------------------
# Structured output parser
# ------------------------------
parser = PydanticOutputParser(pydantic_object=TaskCard)

# ------------------------------
# Prompt template with format instructions
# ------------------------------
prompt = PromptTemplate(
    template=(
        "Ты извлекаешь из неформального описания учебного задания структурированные данные.\n"
        "Верни их в EXACTLY the JSON format described by the format instructions.\n"
        "Описание задания: {task_text}\n"
        "{format_instructions}"
    ),
    input_variables=["task_text"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# ------------------------------
# LLM configuration (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,
)

# ------------------------------
# Chain: prompt -> LLM -> parser
# ------------------------------
chain = prompt | llm | parser

# ------------------------------
# DeepAgents setup (required by the course)
# ------------------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

@tool
def echo_tool(text: str) -> str:
    """Simple echo tool, returns the received text."""
    return text

agent = create_deep_agent(
    model=llm,
    tools=[echo_tool],
    backend=backend,
    system_prompt="You are a helpful assistant that can also run tools.",
)

# ------------------------------
# Main execution
# ------------------------------
async def main():
    # Example input (can be replaced with any other string)
    task_description = (
        "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. "
        "Оценка: за полноту и за пример кода."
    )

    # Use the structured chain to parse the description
    parsed_result: TaskCard = await chain.ainvoke({"task_text": task_description})

    # Print the validated model dump
    print("=== Validated TaskCard ===")
    print(parsed_result.model_dump())

    # Print a short human-readable summary
    print("\n=== Summary ===")
    print(f"Title: {parsed_result.title}")
    print(f"Subject: {parsed_result.subject}")
    print(f"Deadline: {parsed_result.deadline_hint}")
    print(f"Deliverable: {parsed_result.deliverable_type}")
    print(f"Grading hints: {', '.join(parsed_result.grading_hints)}")

    # Demonstrate that the deep agent is functional (optional)
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content="Echo this message")]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    print("\n=== Agent echo result ===")
    print(result["messages"][-1].content)

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