"""
Simple module that converts informal assignment description into a structured data card.

Usage:
    from main import parse_assignment
    card = parse_assignment("Write an essay on climate change by next Friday")
"""
from typing import Dict, Any

from langchain.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

class AssignmentCard(BaseModel):
    title: str = Field(..., description="Short title of the assignment")
    subject: str | None = Field(None, description="Subject or topic of the assignment")
    deadline_hint: str | None = Field(
        None,
        description="Human‑readable hint about when the assignment is due",
    )
    deliverable_type: str | None = Field(
        None,
        description="What kind of work should be submitted (essay, report, code, etc.)",
    )
    grading_hints: str | None = Field(
        None,
        description="Any hints about how the assignment will be graded",
    )

parser = PydanticOutputParser(pydantic_object=AssignmentCard)

# Prompt template that asks the model to output JSON matching AssignmentCard
PROMPT_TEMPLATE = (
    "You are an assistant that converts a short informal description of a study assignment into a structured data card."
    " Return only valid JSON that matches the following schema:\n{schema}\n"
    " Description: {description}"
)

def parse_assignment(description: str) -> AssignmentCard:
    """Return an AssignmentCard parsed from the given description.

    The function uses LangChain's PydanticOutputParser to enforce type safety.
    """
    from langchain import PromptTemplate, LLMChain
    from langchain.chat_models import ChatOpenAI

    # Use a small model for demonstration; replace with your own key if needed.
    llm = ChatOpenAI(temperature=0.2)
    prompt = PromptTemplate(
        input_variables=["description", "schema"],
        template=PROMPT_TEMPLATE,
    )
    chain = LLMChain(llm=llm, prompt=prompt, output_parser=parser)
    result = chain.run(description=description, schema=parser.get_format_instructions())
    return result

if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python main.py '<assignment description>'")
        sys.exit(1)
    desc = sys.argv[1]
    card = parse_assignment(desc)
    print(card.json(indent=4))