"""
Task: Parse a raw assignment description into a flat card.

This script demonstrates a LangChain pipeline that:
1. Defines a Pydantic model for the assignment card.
2. Builds a prompt that asks the model to output JSON conforming to that model.
3. Uses the `PydanticOutputParser` to enforce the structure.
4. Runs the chain on a sample input and prints the validated object and a human‑readable summary.

Requirements:
- langchain-core
- langchain-openai
- pydantic
- python-dotenv (optional, for loading API keys)
"""

import os
from dotenv import load_dotenv

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

# Load environment variables (e.g. OPENAI_API_KEY)
load_dotenv()

# 1. Define the assignment card model
class AssignmentCard(BaseModel):
    """Flat representation of an assignment description.

    Attributes
    ----------
    title : str
        The main title or name of the assignment.
    subject : str
        The academic subject or topic.
    deadline_hint : str
        Free‑form hint about the deadline (e.g. "к пятнице").
    deliverable_type : str
        What the student should submit (e.g. "отчёт", "код").
    grading_hints : list[str]
        List of phrases or criteria mentioned for grading.
    """

    title: str = Field(..., description="Title of the assignment")
    subject: str = Field(..., description="Subject or topic of the assignment")
    deadline_hint: str = Field(..., description="Hint about the deadline")
    deliverable_type: str = Field(..., description="What the student should submit")
    grading_hints: list[str] = Field(..., description="List of grading criteria mentioned")

# 2. Create the parser that will enforce the model
parser = PydanticOutputParser(pydantic_object=AssignmentCard)

# 3. Build the prompt template
prompt_template = PromptTemplate(
    template="""
You are an assistant that extracts structured information from a single sentence or short paragraph describing an assignment.

Given the following description, output a JSON object that conforms exactly to the following schema:
{schema}

The JSON should contain the keys: title, subject, deadline_hint, deliverable_type, grading_hints.

Description:
{description}

Output:
{format_instructions}
""",
    input_variables=["description"],
    partial_variables={
        "schema": parser.get_format_instructions(),
        "format_instructions": parser.get_format_instructions(),
    },
)

# 4. Set up the LLM (ChatOpenAI).  The model name can be overridden via env var.
llm = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
    temperature=0,
)

# 5. Build the chain
chain = prompt_template | llm | parser

# 6. Example usage
if __name__ == "__main__":
    example_description = (
        "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. Оценка: за полноту и за пример кода."
    )
    result = chain.invoke({"description": example_description})
    print("\n--- Parsed Assignment Card ---")
    print(result.model_dump(indent=4))

    # Human‑readable summary
    summary = (
        f"Title: {result.title}\n"
        f"Subject: {result.subject}\n"
        f"Deadline hint: {result.deadline_hint}\n"
        f"Deliverable: {result.deliverable_type}\n"
        f"Grading hints: {', '.join(result.grading_hints)}"
    )
    print("\n--- Summary ---")
    print(summary)
""