import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

# Loading the API Key for ChatGPT
load_dotenv()

while True:
    print("You: (type END on a new line to submit)")

    lines = []

    while True:
        line = input()

        if line.lower() == "exit":
            exit()

        if line.strip() == "END":
            break

        lines.append(line)

    user_input = "\n".join(lines)

    if not user_input.strip():
        continue

    prompt_template = PromptTemplate(
        template="""
You are TextSum, an AI expert that will take the text as input and summarize that text as output.

You have to be very careful about the instructions (if given).
Otherwise, just summarize the text in 5-10 points:

{user_input}
""",
        input_variables=["user_input"]
    )

    llm = ChatOpenAI(
        model="gpt-4o-mini-2024-07-18"
    )

    output_parser = StrOutputParser()

    chain = prompt_template | llm | output_parser

    results = chain.invoke({
        "user_input": user_input
    })

    print(f"\nAI:\n{results}\n")