from typing import TypedDict, Optional

from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langgraph.graph import StateGraph
from pydantic import BaseModel


class TaskCard(BaseModel):
    title: str
    subject: str
    deadline_hint: str
    deliverable_type: str
    grading_hints: str


class State(TypedDict):
    raw_text: str
    task_card: Optional[TaskCard]
    summary: Optional[str]
    error: Optional[str]


# Prompt that instructs the LLM to extract fields in JSON
PROMPT_TEMPLATE = """
You are a data extraction assistant. Extract the following fields from the text provided:
- title
- subject
- deadline_hint
- deliverable_type
- grading_hints

Return ONLY a JSON object with these keys and no additional explanation.

Text:
{raw_text}
"""

prompt = PromptTemplate.from_template(PROMPT_TEMPLATE)

# LLM configuration
llm = ChatOpenAI(model="gpt-4o", temperature=0)

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

# Chain: prompt -> LLM -> parser
from langchain_core.runnables import RunnableParallel, RunnablePassthrough

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


# Node that processes the raw text
async def process_task(state: State) -> State:
    try:
        task_card = await chain.ainvoke({"raw_text": state["raw_text"]})
        summary = (
            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}'."
        )
        state["task_card"] = task_card
        state["summary"] = summary
    except Exception as e:
        state["error"] = str(e)
    return state


# Node that checks for interrupt keyword
async def check_interrupt(state: State) -> State:
    if "interrupt" in state["raw_text"].lower():
        state["error"] = "Interrupt requested. Exiting."
    return state


# Build the StateGraph
graph = StateGraph(State)

# Add nodes
graph.add_node("check_interrupt", check_interrupt)
graph.add_node("process_task", process_task)

# Define edges: start at check_interrupt, then to process_task unless interrupted
graph.set_entry_point("check_interrupt")

# Edge logic: after check_interrupt, if error set, go to END
graph.add_edge("check_interrupt", "process_task", condition=lambda s: s.get("error") is None)
graph.add_edge("check_interrupt", "END", condition=lambda s: s.get("error") is not None)
graph.add_edge("process_task", "END")

# Compile graph
app = graph.compile()


def main():
    raw_text = input("Enter the raw text to parse: ")
    initial_state: State = {
        "raw_text": raw_text,
        "task_card": None,
        "summary": None,
        "error": None,
    }
    final_state = app.invoke(initial_state)

    if final_state["error"]:
        print("Error:", final_state["error"])
    else:
        print("Parsed TaskCard:", final_state["task_card"])
        print("Summary:", final_state["summary"])


if __name__ == "__main__":
    main()