import os
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_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

# Load environment variables
load_dotenv()

# ---------- LLM INITIALIZATION ----------
# Using OpenRouter via langchain-openai as required
llm = ChatOpenAI(
    model_name="gpt-4o-mini",  # example model available on OpenRouter
    base_url="https://openrouter.ai/api/v1/chat/completions",
    api_key=os.getenv("OPENROUTER_API_KEY"),
    temperature=0.0,
)

# ---------- Pydantic MODEL ----------
class TaskCard(BaseModel):
    title: str = Field(description="Краткое название задания")
    subject: str = Field(description="Предмет или тема задания")
    deadline_hint: str = Field(description="Краткая подсказка о сроке сдачи, например 'к пятнице'")
    deliverable_type: str = Field(description="Тип сдаваемого материала: отчёт, код, презентация и т.д.")
    grading_hints: list[str] = Field(description="Список критериев оценки, упомянутых в тексте")

# ---------- PROMPT AND PARSER ----------
parser = PydanticOutputParser(pydantic_object=TaskCard)

prompt_template = """You are an assistant that extracts structured information from a short task description.

Input: {task_text}

Output must be a JSON object with the following fields:
{format_instructions}

Respond ONLY with the JSON object.
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["task_text"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# ---------- TOOL THAT RUNS THE CHAIN ----------
@tool
def parse_task(task_text: str) -> str:
    """Parse a task description into a structured JSON object."""
    chain = prompt | llm | parser
    result = chain.invoke({"task_text": task_text})
    # Return the validated object as JSON string for the agent to forward
    return result.model_dump_json()

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

agent = create_deep_agent(
    model=llm,
    tools=[parse_task],
    backend=backend,
    system_prompt="You are a helpful assistant that parses a single task description into structured data using the provided tool.",
)

# ---------- MAIN EXECUTION ----------
async def main():
    # Example input – single line, no dialogue
    task_description = "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. Оценка: за полноту и за пример кода."
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=task_description)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # The agent will return the tool output as the last message
    output_json = result["messages"][-1].content
    print("\n--- Parsed JSON ---")
    print(output_json)
    # Load into Pydantic for pretty printing and summary
    card = TaskCard.parse_raw(output_json)
    print("\n--- Validated Object ---")
    print(card.model_dump())
    print("\n--- Summary ---")
    print(f"Title: {card.title}\nSubject: {card.subject}\nDeadline: {card.deadline_hint}\nDeliverable: {card.deliverable_type}\nGrading: {', '.join(card.grading_hints)}")

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