from dotenv import load_dotenv
import os

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

# Load environment variables
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise RuntimeError("OPENAI_API_KEY is not set in .env file or environment.")

class TaskCard(BaseModel):
    title: str = Field(..., description="Title of the task")
    subject: str = Field(..., description="Subject of the task")
    deadline_hint: str = Field(..., description="Deadline hint for the task")
    deliverable_type: str = Field(..., description="Type of deliverable")
    grading_hints: str = Field(..., description="Hints for grading")

# Prompt template instructing LLM to output JSON matching TaskCard schema
PROMPT_TEMPLATE = (
    "You are a data extraction assistant. Extract the following fields from the text provided:\n"
    "- title\n"
    "- subject\n"
    "- deadline_hint\n"
    "- deliverable_type\n"
    "- grading_hints\n\n"
    "Return ONLY a JSON object with these keys and no additional explanation.\n\n"
    "Text:\n{raw_text}\n"
)

prompt = PromptTemplate.from_template(PROMPT_TEMPLATE)

# LLM configuration – use the API key loaded from environment
llm = ChatOpenAI(model="gpt-4o", temperature=0, api_key=api_key)

# Parser to enforce TaskCard schema
parser = PydanticOutputParser(pydantic_object=TaskCard)

# Chain: prompt -> llm -> parser
from langchain_core.runnables import RunnablePassthrough

chain = (
    RunnablePassthrough.assign(raw_text=lambda x: x["raw_text"]) | prompt
) | llm | parser


def parse_task(raw_text: str) -> TaskCard:
    """Parse raw text into a validated TaskCard using the LLM chain."""
    try:
        task_card = chain.invoke({"raw_text": raw_text})
        return task_card
    except Exception as e:
        raise RuntimeError(f"Failed to parse task: {e}") from e


def summarize_task(task_card: TaskCard) -> str:
    """Return a human‑readable summary of the task."""
    return (
        f"Task '{task_card.title}' on subject '{task_card.subject}' has a deadline of '{task_card.deadline_hint}'. "
        f"Deliverable type: '{task_card.deliverable_type}'. Grading hints: '{task_card.grading_hints}'."
    )


def main():
    raw_text = input("Enter the raw text to parse: ")
    try:
        task_card = parse_task(raw_text)
    except RuntimeError as e:
        print("Error:", e)
        return
    summary = summarize_task(task_card)
    print("Parsed TaskCard:\n", task_card.json(indent=2))
    print("\nSummary:\n", summary)

if __name__ == "__main__":
    main()