#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Assignment Card Extraction

This script demonstrates how to extract structured assignment details from a
natural language description using LangChain and Pydantic.
"""

import os
from typing import List

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

# Load environment variables (expects OPENAI_API_KEY)
load_dotenv()

# --------------------------------------------------------------------------- #
# Pydantic model definition
# --------------------------------------------------------------------------- #
class AssignmentCard(BaseModel):
    """
    Structured representation of an assignment description.
    """

    title: str = Field(
        ...,
        description="Short title of the assignment (e.g., 'Mini-report on LangChain').",
    )
    subject: str = Field(
        ...,
        description="Subject or topic of the assignment (e.g., 'LangChain').",
    )
    deadline_hint: str = Field(
        ...,
        description="A short phrase indicating the deadline (e.g., 'by Friday').",
    )
    deliverable_type: str = Field(
        ...,
        description="What to submit: report, code, presentation, etc.",
    )
    grading_hints: List[str] = Field(
        ...,
        description="List of key grading criteria mentioned in the description.",
    )

# --------------------------------------------------------------------------- #
# LangChain components
# --------------------------------------------------------------------------- #
# Parser that will convert the LLM output into an AssignmentCard instance
parser = PydanticOutputParser(pydantic_object=AssignmentCard)

# Prompt template that instructs the LLM to output JSON matching the model
prompt = PromptTemplate(
    template=(
        "You are an assignment extraction assistant. "
        "Given the following assignment description, extract the following fields:\n\n"
        "- title: short title of the assignment\n"
        "- subject: subject or topic\n"
        "- deadline_hint: a short phrase indicating the deadline\n"
        "- deliverable_type: what to submit (e.g., report, code, presentation)\n"
        "- grading_hints: list of key grading criteria mentioned\n\n"
        "Return a JSON object with exactly these keys. Do not include any additional keys or text.\n\n"
        "Description: {description}\n\n"
        "{format_instructions}"
    ),
    input_variables=["description"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# LLM configuration
llm = ChatOpenAI(
    temperature=0,
    model="gpt-3.5-turbo",
)

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

# --------------------------------------------------------------------------- #
# Main execution
# --------------------------------------------------------------------------- #
def main() -> None:
    # Sample assignment description
    sample_description = (
        "Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. "
        "Оценка: за полноту и за пример кода."
    )

    # Run the chain
    try:
        result = chain.invoke({"description": sample_description})
    except Exception as e:
        print(f"Error during chain execution: {e}")
        return

    # The result is already a validated AssignmentCard instance
    print("\n=== Parsed Assignment Card ===")
    print(result.model_dump(indent=2))

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


if __name__ == "__main__":
    main()