import gradio as gr
import random
import os
import re
from huggingface_hub import InferenceClient
from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound
from fpdf import FPDF
from fpdf.enums import XPos, YPos
from datetime import datetime

# 클라이언트 생성 함수
def create_client(model_name):
    return InferenceClient(model_name, token=os.getenv("HF_TOKEN"))

client = create_client("CohereForAI/c4ai-command-r-plus")

# API 호출 함수
def call_api(content, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
    random_seed = random.randint(0, 1000000)
    response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
    return response.choices[0].message.content

# 정보 분석 함수
def analyze_info(category, style, transcripts):
    transcript_list = transcripts.split("\n\n---\n\n")
    analyzed_content = f"선택한 카테고리: {category}\n선택한 포스팅 스타일: {style}\n\n"
    for i, transcript in enumerate(transcript_list, 1):
        analyzed_content += f"유튜브 트랜스크립트 {i}:\n{transcript}\n\n"
    return analyzed_content

# 블로그 포스트 생성 함수
def generate_blog_post(category, style, transcripts, category_prompt, style_prompt, max_tokens, temperature, top_p):
    full_content = analyze_info(category, style, transcripts)
    combined_prompt = f"{category_prompt}\n\n{style_prompt}\n\n{full_content}"
    modified_text = call_api(combined_prompt, "", max_tokens, temperature, top_p)
    return modified_text.replace('\n', '\n\n')

# 유튜브 대본 요약 함수
def summarize_transcript(transcripts, system_message, max_tokens, temperature, top_p):
    summary = call_api(transcripts, system_message, max_tokens, temperature, top_p)
    return summary

# 유튜브 비디오 ID 추출 함수
def get_video_id(youtube_url):
    video_id_match = re.search(r"(?<=v=)[^#&?]*", youtube_url) or re.search(r"(?<=youtu.be/)[^#&?]*", youtube_url)
    return video_id_match.group(0) if video_id_match else None

# 유튜브 트랜스크립트 추출 함수
def get_transcript(youtube_url):
    video_id = get_video_id(youtube_url)
    if not video_id:
        return "Invalid YouTube URL. Please enter a valid URL."

    language_order = ['ko', 'en', 'ja', 'zh-Hans', 'pt', 'es', 'it', 'fr', 'de', 'ru']
    for lang in language_order:
        try:
            transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=[lang])
            return f"Transcript in {lang}:\n\n{' '.join([entry['text'] for entry in transcript])}"
        except NoTranscriptFound:
            continue
        except Exception as e:
            return f"Error: {str(e)}"
    return "No transcript available in the specified languages."

# 카테고리별 프롬프트 함수
def get_blog_post_prompt(category):
    if category == "일반형":
        return """
        #유튜브 대본을 블로그 포스팅으로 변환하는 규칙(일반형_v4)
        
        """
    elif category == "정보성":
        return """
        #유튜브 대본을 블로그 포스팅으로 변환하는 규칙(정보성_v4)
        
        """
    elif category == "1개 상품 추천형":
        return """
        #유튜브 대본을 블로그 포스팅으로 변환하는 규칙(추천형_v4)
       
        """
    elif category == "큐레이션형":
        return """
        #유튜브 대본을 블로그 포스팅으로 변환하는 규칙(큐레이션형_v3)
        
        """        

# 포스팅 스타일 프롬프트 함수
def get_style_prompt(style):
    prompts = {
        "친근한": """
 제가 알게 된 꿀팁들을 하나하나 알려드릴게요. 
""",
        "일반":"""#일반적인 블로그 포스팅 스타일 가이드

""",
        "전문적인": """
#전문적인 블로그 포스팅 스타일 가이드

"""
    }
    return prompts.get(style, "포스팅 스타일 프롬프트")

# 포스팅 스타일 설명 함수
def get_style_description(style):
    descriptions = {
        "친근한": "독자와 가까운 친구처럼 대화하는 듯한 친근한 스타일입니다.",
        "일반": "일반적이고 중립적인 톤으로 정보를 전달하는 스타일입니다.",
        "전문적인": "전문가의 시각에서 깊이 있는 정보를 전달하는 스타일입니다."
    }
    return descriptions.get(style, "포스팅 스타일을 선택하세요.")

# 프롬프트 업데이트 함수
def update_prompts_and_description(category, style):
    blog_post_prompt = get_blog_post_prompt(category)
    style_prompt = get_style_prompt(style)
    style_description = get_style_description(style)
    return blog_post_prompt, style_prompt, style_description

def format_filename(text):
    text = re.sub(r'[^\w\s-]', '', text)
    return text[:50].strip()

def extract_first_recommended_title(blog_post):
    section_match = re.search(r'(?:#+\s*)?추천\s*제목:?\s*\n([\s\S]*?)(?=\n(?:#+|$)|$)', blog_post, re.IGNORECASE)
    if section_match:
        section_content = section_match.group(1)
        title_match = re.search(r'(?:^|\n)\s*(?:\d+\.|-|\*|\•)?\s*(.*?)(?=\n|$)', section_content)
        if title_match:
            title = title_match.group(1).strip()
            print(f"Extracted title: {title}")
            return title
    print("No title found")
    return "블로그_글"

class PDF(FPDF):
    def __init__(self):
        super().__init__()
        self.add_font("NanumGothic", "", "NanumGothic.ttf")
        self.add_font("NanumGothicBold", "", "NanumGothicBold.ttf")
        self.add_font("NanumGothicExtraBold", "", "NanumGothicExtraBold.ttf")
        self.add_font("NanumGothicLight", "", "NanumGothicLight.ttf")

    def header(self):
        # 헤더를 비워둡니다
        pass

    def footer(self):
        self.set_y(-15)
        self.set_font('NanumGothicLight', '', 8)
        self.cell(0, 10, f'Page {self.page_no()}', 0, new_x=XPos.RIGHT, new_y=YPos.TOP, align='C')

def save_to_pdf(summary, blog_post, file_type):
    pdf = PDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    
    pdf.add_page()
    pdf.set_font("NanumGothicExtraBold", size=16)
    pdf.cell(0, 10, "요약", new_x=XPos.LMARGIN, new_y=YPos.NEXT, align='C')
    pdf.ln(5)
    pdf.set_font("NanumGothic", size=11)
    pdf.multi_cell(0, 6, summary)

    pdf.add_page()
    pdf.set_font("NanumGothicExtraBold", size=16)
    pdf.cell(0, 10, "블로그 글", new_x=XPos.LMARGIN, new_y=YPos.NEXT, align='C')
    pdf.ln(5)

    lines = blog_post.split('\n')
    for line in lines:
        if line.strip() == '':
            pdf.ln(3)  # 빈 줄은 작은 간격만 추가
        elif line.startswith('#'):  # 제목으로 간주
            pdf.set_font("NanumGothicBold", size=14)
            pdf.multi_cell(0, 8, line.lstrip('#').strip())
            pdf.ln(2)
        elif line.startswith('##'):  # 부제목으로 간주
            pdf.set_font("NanumGothicBold", size=12)
            pdf.multi_cell(0, 7, line.lstrip('#').strip())
            pdf.ln(2)
        else:
            pdf.set_font("NanumGothic", size=11)
            pdf.multi_cell(0, 6, line.strip())
            pdf.ln(1)

    title = extract_first_recommended_title(blog_post)
    today_date = datetime.now().strftime("%Y%m%d")
    filename = f"{today_date}_{format_filename(title)}.pdf"
    
    print(f"Saving PDF as: {filename}")
    pdf.output(filename)
    return filename

# Gradio 인터페이스용 PDF 저장 함수
def save_content_to_pdf(summary, blog_post):
    filename = save_to_pdf(summary, blog_post, "블로그")
    return filename

# Gradio 인터페이스 구성
title = "유튜브로 블로그 글 생성하기"

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}")
    
    # 1단계: 카테고리 선택
    gr.Markdown("### 1단계: 포스팅 카테고리를 지정해주세요", elem_id="step-title")
    category = gr.Radio(choices=["일반형","정보성", "1개 상품 추천형", "큐레이션형"], label="포스팅 카테고리", value="일반형")
    
    # 구분선 추가
    gr.Markdown("---\n\n")

    # 2단계: 포스팅 스타일 선택
    gr.Markdown("### 2단계: 포스팅 스타일을 선택해주세요", elem_id="step-title")
    style = gr.Radio(choices=["친근한", "일반", "전문적인"], label="포스팅 스타일", value="친근한")
    style_description = gr.Markdown(f"{get_style_description('친근한')}", elem_id="style-description")
    
    # 구분선 추가
    gr.Markdown("---\n\n")

    # 3단계: 유튜브 링크를 입력하세요
    gr.Markdown("### 3단계: 유튜브 링크를 입력하세요", elem_id="step-title")
    with gr.Row():
        youtube_url1 = gr.Textbox(label="YouTube URL 1", placeholder="첫 번째 유튜브 링크를 입력하세요")
        youtube_url2 = gr.Textbox(label="YouTube URL 2", placeholder="두 번째 유튜브 링크를 입력하세요")
        youtube_url3 = gr.Textbox(label="YouTube URL 3", placeholder="세 번째 유튜브 링크를 입력하세요")
    
    # 숨겨진 텍스트박스 (사용자에게 보이지 않음)
    combined_urls = gr.Textbox(visible=False)
    transcript_output = gr.Textbox(label="유튜브 트랜스크립트", lines=10)
    
    # 유튜브 트랜스크립트 가져오기 함수
    def combine_and_get_transcripts(url1, url2, url3):
        urls = [url for url in [url1, url2, url3] if url.strip()]
        combined = ",".join(urls)
        all_transcripts = []
        for url in urls:
            transcript = get_transcript(url.strip())
            all_transcripts.append(transcript)
        return combined, "\n\n---\n\n".join(all_transcripts)
    
    # 입력 변경 시 트랜스크립트 업데이트
    for url_input in [youtube_url1, youtube_url2, youtube_url3]:
        url_input.change(
            fn=combine_and_get_transcripts, 
            inputs=[youtube_url1, youtube_url2, youtube_url3], 
            outputs=[combined_urls, transcript_output]
        )
    
    # 요약글 생성하기 기능 추가
    gr.Markdown("### 요약글 생성하기", elem_id="step-title")
    with gr.Accordion("요약글 설정", open=False):
        summary_system_message = gr.Textbox(
            label="요약글 시스템 메시지", 
            value="""
        #유튜브 대본 요약 규칙
        
            """, 
            lines=15, 
            visible=True
        )
        summary_max_tokens = gr.Slider(label="Max Tokens", minimum=1000, maximum=7000, value=5000, step=1000)
        summary_temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.05)
        summary_top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
    
    summarize_btn = gr.Button("요약글 생성하기")
    summary_output = gr.Textbox(label="요약된 글", lines=10)
    
    def generate_summary(transcript, system_message, max_tokens, temperature, top_p):
        summary = summarize_transcript(transcript, system_message, max_tokens, temperature, top_p)
        return summary

    summarize_btn.click(
        fn=generate_summary, 
        inputs=[transcript_output, summary_system_message, summary_max_tokens, summary_temperature, summary_top_p], 
        outputs=[summary_output]
    )

    # 구분선 추가
    gr.Markdown("---\n\n")

    # 4단계: 글 생성하기
    gr.Markdown("### 4단계: 글 생성하기", elem_id="step-title")
    gr.HTML("<span style='color: grey;'>[생성하기 버튼을 선택해주세요]</span>")
    
    with gr.Accordion("블로그 글 설정", open=False):
        blog_system_message = gr.Textbox(label="카테고리 프롬프트", value=get_blog_post_prompt("일반형"), lines=20, visible=True)
        style_prompt_hidden = gr.Textbox(label="스타일 프롬프트", value=get_style_prompt("친근한"), lines=10, visible=False)  # 초기값 설정
        blog_max_tokens = gr.Slider(label="Max Tokens", minimum=1000, maximum=12000, value=8000, step=1000)
        blog_temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.8, step=0.1)
        blog_top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
        
    generate_btn = gr.Button("블로그 글 생성하기")
    blog_output = gr.Textbox(label="생성된 블로그 글", lines=30)
    
    def generate_blog_content(category, style, transcripts, category_prompt, style_prompt, max_tokens, temperature, top_p):
        blog_post = generate_blog_post(category, style, transcripts, category_prompt, style_prompt, max_tokens, temperature, top_p)
        return blog_post

    generate_btn.click(
    fn=generate_blog_content, 
    inputs=[category, style, transcript_output, blog_system_message, style_prompt_hidden, blog_max_tokens, blog_temperature, blog_top_p], 
    outputs=[blog_output]
)

    # PDF 저장 버튼 추가
    save_pdf_btn = gr.Button("PDF로 저장하기")
    pdf_output = gr.File(label="생성된 PDF 파일")

    save_pdf_btn.click(
        fn=save_content_to_pdf,
        inputs=[summary_output, blog_output],
        outputs=[pdf_output]
    )

    # 카테고리와 스타일이 변경될 때 프롬프트 업데이트
    def update_prompts_and_description(category, style):
        blog_post_prompt = get_blog_post_prompt(category)
        style_prompt = get_style_prompt(style)
        style_description = get_style_description(style)
        return blog_post_prompt, style_prompt, style_description
    
    category.change(fn=update_prompts_and_description, inputs=[category, style], outputs=[blog_system_message, style_prompt_hidden, style_description])
    style.change(fn=update_prompts_and_description, inputs=[category, style], outputs=[blog_system_message, style_prompt_hidden, style_description])

demo.launch()

# CSS 스타일 추가
gr.HTML("""
<style>
#step-title {
    font-size: 1.7em;
    font-weight: bold;
}
#style-description {
    font-size: 1.2em;
}
</style>
""")