import gradio as gr
import requests
import os
from openai import OpenAI

UPSTAGE_API_KEY = os.getenv("UPSTAGE_API_KEY")

def parse_document(file):
    url = "https://api.upstage.ai/v1/document-ai/document-parse"
    headers = {'Authorization': f'Bearer {UPSTAGE_API_KEY}'}
    files = {"document": open(file.name, "rb")}
    data = {
        "base64_encoding": "['table']",
        "model": "document-parse"
    }

    response = requests.post(url, headers=headers, files=files, data=data)
    result = response.json()
    html_text = result.get("content", {}).get("html", "")
    return html_text

def chat_with_document(history, html_text, user_question):
    if not html_text.strip():
        return history, history, "⚠️ 먼저 문서를 변환해주세요."

    client = OpenAI(
        api_key=UPSTAGE_API_KEY,
        base_url="https://api.upstage.ai/v1"
    )

    history = history or []

    system_prompt = f"""The following is a financial statement document extracted in HTML format.
    Please answer user questions accurately and concisely in Korean, based on the text within HTML tags.

    Document:
    {html_text}
    """

    messages = [{"role": "system", "content": system_prompt}]
    for user, bot in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": user_question})

    try:
        response = client.chat.completions.create(
            model="solar-pro",
            messages=messages,
            temperature=0,
            max_tokens=1024
        )
        bot_reply = response.choices[0].message.content
    except Exception as e:
        bot_reply = f"⚠️ 오류가 발생했습니다: {str(e)}"

    history.append((user_question, bot_reply))
    return history, history, ""

def toggle_html_view(current_html, is_visible):
    return (
        gr.update(value=current_html, visible=not is_visible),
        gr.update(value=current_html, visible=is_visible),
        not is_visible
    )

with gr.Blocks() as demo:

    gr.Markdown("# 📄 재무제표 분석 챗봇")
    gr.Markdown("1. Document Parse API로 PDF 문서를 HTML로 변환합니다.\n"
                "2. Solar LLM을 통해 문서 기반 질문에 답변합니다.")
    gr.Markdown("예제 파일은 Files 버튼을 클릭하면 확인 및 다운로드 가능합니다.")

    with gr.Row():
        file_input = gr.File(label="📎 재무제표 업로드")
        parse_btn = gr.Button("문서 HTML 변환")

    html_output = gr.Textbox(label="📘 문서 내용", lines=10, visible=True, elem_id="scrollable-html")
    html_display = gr.HTML(visible=False, elem_id="scrollable-html-display")
    toggle_html_btn = gr.Button("🔁 HTML 보기 전환")
    html_visible_state = gr.State(False)

    parse_btn.click(fn=parse_document, inputs=file_input, outputs=html_output)
    toggle_html_btn.click(
        fn=toggle_html_view,
        inputs=[html_output, html_visible_state],
        outputs=[html_output, html_display, html_visible_state]
    )

    chatbot = gr.Chatbot(label="💬 문서 기반 Q&A", height=400)
    user_question = gr.Textbox(label="❓ 질문을 입력하세요", lines=2)
    answer_btn = gr.Button("답변 생성")
    chat_state = gr.State([])

    with gr.Row():
        gr.Markdown("💡 예제 질문:")
        ex1 = gr.Button("어떤 기업의 재무제표인가요?")
        ex2 = gr.Button("Q3 분기의 총 매출액은 얼마인가요?")

    # 예제 질문 → 질문 입력 + 자동 응답
    ex1.click(
        fn=lambda: "어떤 기업의 재무제표인가요?",
        inputs=[],
        outputs=user_question
    ).then(
        fn=chat_with_document,
        inputs=[chat_state, html_output, user_question],
        outputs=[chatbot, chat_state, user_question],
        show_progress=True
    )

    ex2.click(
        fn=lambda: "Q3 분기의 총 매출액은 얼마인가요?",
        inputs=[],
        outputs=user_question
    ).then(
        fn=chat_with_document,
        inputs=[chat_state, html_output, user_question],
        outputs=[chatbot, chat_state, user_question],
        show_progress=True
    )

    answer_btn.click(
        fn=chat_with_document,
        inputs=[chat_state, html_output, user_question],
        outputs=[chatbot, chat_state, user_question],
        show_progress=True
    )

demo.css = """
#scrollable-html, #scrollable-html-display {
    max-height: 400px;
    overflow: auto;
    border: 1px solid #ccc;
    padding: 10px;
}
"""

demo.launch()