import torch
import gradio as gr
import requests
from bs4 import BeautifulSoup
from transformers import T5Tokenizer, T5ForConditionalGeneration, T5TokenizerFast, AutoTokenizer, AutoModelForSeq2SeqLM
import nltk

# nltk의 'punkt' 데이터 다운로드 (문장 토큰화를 위해)
nltk.download('punkt')

# 뉴스 기사 본문 또는 URL을 입력받아,
# 만약 입력값이 "http"로 시작하면 네이버 뉴스 URL로 간주하여 본문을 추출하고,
# 그렇지 않으면 그대로 텍스트를 반환하는 함수
def extract_naver_news_article(text_or_url):
    try:
        text_or_url = text_or_url.strip()
        if text_or_url.startswith("http"):
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
            }
            response = requests.get(text_or_url, headers=headers)
            if response.status_code != 200:
                return f"페이지 요청 실패 (HTTP {response.status_code})"
            soup = BeautifulSoup(response.text, "html.parser")
            article_content = soup.select_one("#dic_area")
            if not article_content:
                article_content = soup.select_one(".go_trans._article_content")
            if not article_content:
                return "기사 본문을 찾을 수 없습니다."
            paragraphs = article_content.find_all("p")
            article_text = "\n".join([p.get_text(strip=True) for p in paragraphs]) if paragraphs else article_content.get_text(strip=True)
            return article_text
        else:
            return text_or_url
    except Exception as e:
        return f"오류 발생: {e}"

# 제목 생성 모델 로드 함수 (파인튜닝된 T5 모델)
def load_title_model(checkpoint_path):
    try:
        title_tokenizer = T5TokenizerFast.from_pretrained(checkpoint_path)
        title_model = T5ForConditionalGeneration.from_pretrained(checkpoint_path)
        return title_model, title_tokenizer
    except Exception as e:
        print(f"Error loading title model: {e}")
        return None, None

# 뉴스 기사 요약 모델 로드 함수 
def load_summarization_model(model_dir):
    try:
        summarizer_tokenizer = AutoTokenizer.from_pretrained(model_dir)
        summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
        return summarizer_model, summarizer_tokenizer
    except Exception as e:
        print(f"Error loading summarization model: {e}")
        return None, None

# 모델 경로 설정 (사용자 환경에 맞게 변경)
title_path = 'onebeans/keT5-news-title-gen'  # 파인튜닝된 T5 제목 생성 모델 경로
abs_path =  'onebeans/keT5-news-summarizer'                 # 뉴스 기사 요약 모델 경로

title_model, title_tokenizer = load_title_model(title_path)
summarizer_model, summarizer_tokenizer = load_summarization_model(abs_path)

# 뉴스 기사 제목 생성 함수
def generate_title(article_text, max_length=20, num_beams=10, early_stopping=True):
    try:
        input_ids = title_tokenizer.encode(f"summarize: {article_text}", return_tensors="pt", truncation=True)
        outputs = title_model.generate(
            input_ids,
            max_length=max_length,
            num_beams=num_beams,
            early_stopping=early_stopping
        )
        title = title_tokenizer.decode(outputs[0], skip_special_tokens=True)
        return title
    except Exception as e:
        return f"Error generating title: {e}"

# 뉴스 기사 요약 생성 함수
def generate_summary(article_text, max_length=40, num_beams=10, early_stopping=True):
    try:
        input_ids = summarizer_tokenizer.encode(article_text, return_tensors="pt", truncation=True)
        outputs = summarizer_model.generate(
            input_ids,
            max_length=max_length,
            num_beams=num_beams,
            early_stopping=early_stopping
        )
        summary = summarizer_tokenizer.decode(outputs[0], skip_special_tokens=True)
        return summary
    except Exception as e:
        return f"Error generating summary: {e}"

# 뉴스 기사 본문 또는 URL을 입력받아 본문 추출 후 제목과 요약 생성하는 함수
def process_news(input_text, title_max_length, title_num_beams, title_early_stopping,
                 summary_max_length, summary_num_beams, summary_early_stopping):
    article_text = extract_naver_news_article(input_text)
    title = generate_title(article_text, max_length=title_max_length, num_beams=title_num_beams, early_stopping=title_early_stopping)
    summary = generate_summary(article_text, max_length=summary_max_length, num_beams=summary_num_beams, early_stopping=summary_early_stopping)
    return article_text, title, summary

# Gradio 인터페이스 실행 함수
def launch_gradio_interface():
    interface = gr.Interface(
        fn=process_news,
        inputs=[
            gr.Textbox(label="뉴스 기사 본문 또는 URL", placeholder="뉴스 기사 본문을 직접 입력하거나 URL을 입력하세요.", lines=10),
            gr.Slider(0, 40, value=20, step=1, label="제목 생성 - 최대 길이"),
            gr.Slider(1, 10, value=10, step=1, label="제목 생성 - 탐색 빔 수"),
            gr.Checkbox(value=True, label="제목 생성 - 조기 종료"),
            gr.Slider(0, 256, value=128, step=1, label="요약 생성 - 최대 길이"),
            gr.Slider(1, 10, value=10, step=1, label="요약 생성 - 탐색 빔 수"),
            gr.Checkbox(value=True, label="요약 생성 - 조기 종료")
        ],
        outputs=[
            gr.Textbox(label="추출되거나 입력된 기사 본문"),
            gr.Textbox(label="생성된 제목"),
            gr.Textbox(label="생성된 요약문")
        ],
        title="네이버 뉴스 제목 및 요약 생성기",
        description="뉴스 기사 본문 또는 URL을 입력하면, 본문을 자동 추출한 후 제목과 요약문을 생성합니다."
    )
    interface.launch(share=True)

if __name__ == '__main__':
    launch_gradio_interface()