import os
import asyncio
from typing import List

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.messages import HumanMessage
from langchain.tools import tool

from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# -------------------- Pydantic model --------------------
class AssignmentCard(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=AssignmentCard)

# -------------------- Prompt template --------------------
prompt_template = PromptTemplate(
    template=(
        "Ты получаешь неформальное описание учебного задания и должен вернуть "
        "структурированные данные в формате JSON, соответствующем схеме ниже.\n"
        "Верни только JSON, без пояснений.\n"
        "Описание задания:\n{task_description}\n\n"
        "{format_instructions}"
    ),
    input_variables=["task_description"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# -------------------- 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,
)

# -------------------- Tool that runs the chain --------------------
@tool
def parse_assignment(task_description: str) -> str:
    """
    Parse a free-form assignment description into a structured AssignmentCard.
    Returns a JSON string that can be parsed by the Pydantic model.
    """
    chain = prompt_template | llm | parser
    result = chain.invoke({"task_description": task_description})
    # result is an AssignmentCard instance
    return result.model_dump_json(indent=2)

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

agent = create_deep_agent(
    model=llm,
    tools=[parse_assignment],
    backend=backend,
    system_prompt="You are a helpful assistant that extracts assignment data.",
)

# -------------------- Main execution --------------------
async def main():
    # Пример входного текста
    user_input = (
        "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. "
        "Оценка: за полноту и за пример кода."
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_input)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # Последнее сообщение агента содержит JSON строку
    json_output = result["messages"][-1].content
    print("=== Structured JSON ===")
    print(json_output)

    # Для наглядности выводим человекочитаемую сводку
    try:
        import json

        data = json.loads(json_output)
        print("\n=== Human-readable summary ===")
        print(f"Title: {data.get('title')}")
        print(f"Subject: {data.get('subject')}")
        print(f"Deadline hint: {data.get('deadline_hint')}")
        print(f"Deliverable type: {data.get('deliverable_type')}")
        print(f"Grading hints: {', '.join(data.get('grading_hints', []))}")
    except Exception as e:
        print("Failed to parse JSON:", e)


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