from langchain_ollama import ChatOllama
from langchain_ibm import ChatWatsonx
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import (
    StrOutputParser,
    JsonOutputParser,
    PydanticOutputParser,
)
from pydantic import BaseModel, Field
from typing import Literal
from dotenv import load_dotenv
import os
import gradio as gr

qwen_llm = ChatOllama(model="qwen3.5:4b")
exaone_llm = ChatOllama(model="exaone3.5:2.4b")

system_prompt = """\
당신은 뉴스 분석 전문가입니다. 반드시 json 형식으로만 응답하세요.
다른 텍스트 없이 아래 형식으로만 반환하세요\n
{{"title":"기사제목", "date":"작성일자", "keywords":["키워드1", "키워드2", "키워드3"], "category" : "카테고리"}}
"""

template = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{article}"),
    ]
)

chain = template | exaone_llm | JsonOutputParser()


# def news_input(article1, article2, article3):
# # 3개의 기사를 리스트로 묶음
# articles = [article1, article2, article3]
# results = []

# # 각 기사별로 체인 실행
# for text in articles:
#     if text.strip():  # 내용이 있을 때만 실행
#         res = chain.invoke({"article": text})
#         results.append(res)

# # 결과를 하나로 합쳐서 반환
# return "\n\n---\n\n".join(results)


def news_input(texts):
    # ===을 기준으로 기사 분리
    articles = [article for article in texts.split("===") if article.strip()]

    response = chain.batch([{"article": article} for article in articles])
    return "\n\n".join(str(item) for item in response)


app = gr.Interface(
    news_input,
    # inputs=[
    #     gr.Textbox(label="기사 1", lines=5),
    #     gr.Textbox(label="기사 2", lines=5),
    #     gr.Textbox(label="기사 3", lines=5),
    # ],
    inputs=[
        gr.Textbox(
            label="기사",
            placeholder="여러 기사 입력 시 구분자로 ===을 사용하세요",
            lines=20,
        ),
    ],
    outputs=[
        gr.Textbox(label="요약", lines=20),
    ],
    title="✒ 뉴스 분석 전문가",
    description="뉴스 기사에서 정보를 추출합니다.",
)

app.launch()