from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from typing import List

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="Short hint about deadline")
    deliverable_type: str | None = Field(None, description="Type of deliverable (e.g., code, report)")
    grading_hints: List[str] | None = Field(None, description="Hints for grading")

prompt_template = ChatPromptTemplate.from_messages([
    (
        "system",
        "You are an assistant that extracts structured task information from a raw text. Return JSON matching the TaskCard schema.",
    ),
    ("human", "{raw_text}"),
])

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

async def parse_task(raw_text: str) -> TaskCard:
    chain = prompt_template | llm.with_structured_output(TaskCard)
    result = await chain.ainvoke({"raw_text": raw_text})
    return result

if __name__ == "__main__":
    import sys, json
    if len(sys.argv) < 2:
        print("Usage: python main.py <path_to_raw_task_file>")
        sys.exit(1)
    with open(sys.argv[1], "r", encoding="utf-8") as f:
        raw = f.read()
    card = asyncio.run(parse_task(raw))
    print(json.dumps(card.model_dump(), indent=2))