"""
Task: Convert raw assignment text into a flat card using LangChain.

The script defines a Pydantic model `TaskCard` and uses LangChain to prompt an LLM
to output the fields in a JSON format that can be parsed by
`PydanticOutputParser`.  The result is printed as a validated object and a short
summary.
"""

from __future__ import annotations

import json
from typing import Any, Dict

from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# 1. Define the data model for a task card.
# ---------------------------------------------------------------------------
class TaskCard(BaseModel):
    title: str = Field(..., description="Task title")
    subject: str | None = Field(None, description="Subject or topic of the task")
    deadline_hint: str | None = Field(
        None,
        description="Human‑readable hint about when the task should be finished",
    )
    deliverable_type: str | None = Field(
        None,
        description="What kind of output is expected (e.g., code, report)",
    )
    grading_hints: str | None = Field(
        None,
        description="Hints for how the task will be graded",
    )

# ---------------------------------------------------------------------------
# 2. Prompt template – we ask the model to return a JSON object that matches
#    TaskCard.
# ---------------------------------------------------------------------------
prompt_template = (
    "You are an assistant that extracts structured information from a raw text.
    Return only a JSON object with the following keys: title, subject,
    deadline_hint, deliverable_type, grading_hints.  Do not add any
    surrounding text or comments.

    Raw text:
    {raw_text}
"
)
prompt = PromptTemplate.from_template(prompt_template)

# ---------------------------------------------------------------------------
# 3. LLM chain – use OpenAI chat model via LangChain.
# ---------------------------------------------------------------------------
llm = ChatOpenAI(temperature=0, model_name="gpt-4o-mini")
parser = PydanticOutputParser(pydantic_object=TaskCard)

# The chain: prompt -> LLM -> parser
from langchain.chains import LLMChain
chain = LLMChain(llm=llm, prompt=prompt, output_parser=parser)

# ---------------------------------------------------------------------------
# 4. Example usage – replace RAW_TEXT with the assignment description.
# ---------------------------------------------------------------------------
RAW_TEXT = """
## Цель
Научиться из одного пользовательского текста получить проверяемый набор полей (title, subject, deadline_hint, deliverable_type, grading_hints) без диалога и без «ручного» разбора строки в Python.

## Стек
- Python 3.10+
- langchain-core, langchain-openai, pydantic
- PydanticOutputParser для структурированного вывода

## Что нужно сделать
1. Описать Pydantic-модель карточки задания
2. Собрать цепочку: шаблон промпта → вызов LLM → парсер в BaseModel
3. В промпте попросить модель вернуть данные в формате для парсера
4. Вывести валидированный объект и краткую сводку
"""

if __name__ == "__main__":
    # Run the chain
    result: TaskCard = chain.run(raw_text=RAW_TEXT)
    print("Validated TaskCard:\n", result.json(indent=2))
    # Simple summary
    summary = (
        f"Title: {result.title}\n"
        f"Subject: {result.subject or 'N/A'}\n"
        f"Deadline hint: {result.deadline_hint or 'N/A'}\n"
        f"Deliverable type: {result.deliverable_type or 'N/A'}\n"
        f"Grading hints: {result.grading_hints or 'N/A'}"
    )
    print("\nSummary:\n", summary)

# End of file