import gradio as gr
from huggingface_hub import InferenceClient
import os
import pandas as pd
from typing import List, Dict, Tuple
import json
import io

import traceback
# 추론 API 클라이언트 설정
hf_client = InferenceClient(
    "CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN")
)

def load_code(filename: str) -> str:
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            return file.read()
    except FileNotFoundError:
        return f"{filename} 파일을 찾을 수 없습니다."
    except Exception as e:
        return f"파일을 읽는 중 오류가 발생했습니다: {str(e)}"

def load_parquet(filename: str) -> str:
    try:
        df = pd.read_parquet(filename, engine='pyarrow')
        return df.head(10).to_markdown(index=False)
    except FileNotFoundError:
        return f"{filename} 파일을 찾을 수 없습니다."
    except Exception as e:
        return f"파일을 읽는 중 오류가 발생했습니다: {str(e)}"


def respond(
    message: str,
    history: List[Dict[str, str]],
    system_message: str = "",
    max_tokens: int = 4000,
    temperature: float = 0.5,
    top_p: float = 0.9,
    parquet_data: str = None
) -> str:
    # 시스템 프롬프트 설정
    if parquet_data:
        system_prefix = """반드시 한글로 답변할 것. 너는 업로드된 데이터를 기반으로 질문에 답변하는 역할을 한다. 데이터를 분석하여 사용자에게 도움이 되는 정보를 제공하라. 데이터를 활용하여 상세하고 정확한 답변을 제공하되, 민감한 정보나 개인 정보를 노출하지 마라."""
        try:
            df = pd.read_json(io.StringIO(parquet_data))
            # 데이터의 요약 정보 생성
            data_summary = df.describe(include='all').to_string()
            system_prefix += f"\n\n업로드된 데이터의 요약 정보:\n{data_summary}"
        except Exception as e:
            print(f"데이터 로드 중 오류 발생: {str(e)}\n{traceback.format_exc()}")
            system_prefix += "\n\n데이터를 로드하는 중 오류가 발생했습니다."
    else:
        system_prefix = system_message or "너는 AI 조언자 역할이다."

    # 메시지 생성
    prompt = system_prefix + "\n\n"
    for chat in history:
        if chat['role'] == 'user':
            prompt += f"사용자: {chat['content']}\n"
        else:
            prompt += f"AI: {chat['content']}\n"
    prompt += f"사용자: {message}\nAI:"

    try:
        # 모델에 메시지 전송 및 응답 받기
        response = ""
        stream = hf_client.text_generation(
            prompt=prompt,
            max_new_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )
        for msg in stream:
            if msg:
                response += msg
                yield response
    except Exception as e:
        error_message = f"추론 중 오류가 발생했습니다: {str(e)}\n{traceback.format_exc()}"
        print(error_message)
        yield error_message


def upload_csv(file_path: str) -> Tuple[str, str]:
    try:
        # CSV 파일 읽기
        df = pd.read_csv(file_path, sep=',')
        # 필수 컬럼 확인
        required_columns = {'id', 'text', 'label', 'metadata'}
        available_columns = set(df.columns)
        missing_columns = required_columns - available_columns
        if missing_columns:
            return f"CSV 파일에 다음 필수 컬럼이 누락되었습니다: {', '.join(missing_columns)}", ""
        # 데이터 클렌징
        df.drop_duplicates(inplace=True)
        df.fillna('', inplace=True)
        # 데이터 유형 최적화
        df = df.astype({'id': 'int32', 'text': 'string', 'label': 'category', 'metadata': 'string'})
        # Parquet 파일로 변환
        parquet_filename = os.path.splitext(os.path.basename(file_path))[0] + '.parquet'
        df.to_parquet(parquet_filename, engine='pyarrow', compression='snappy')
        return f"{parquet_filename} 파일이 성공적으로 업로드되고 변환되었습니다.", parquet_filename
    except Exception as e:
        return f"CSV 파일 업로드 및 변환 중 오류가 발생했습니다: {str(e)}", ""

def upload_parquet(file_path: str) -> Tuple[str, str, str]:
    try:
        # Parquet 파일 읽기
        df = pd.read_parquet(file_path, engine='pyarrow')
        # Markdown으로 변환하여 미리보기
        parquet_content = df.head(10).to_markdown(index=False)
        # DataFrame을 JSON 문자열로 변환
        parquet_json = df.to_json(orient='records', force_ascii=False)
        return "Parquet 파일이 성공적으로 업로드되었습니다.", parquet_content, parquet_json
    except Exception as e:
        return f"Parquet 파일 업로드 중 오류가 발생했습니다: {str(e)}", "", ""

def text_to_parquet(text: str) -> Tuple[str, str, str]:
    try:
        # 텍스트를 DataFrame으로 변환 (각 행은 콤마로 구분)
        data = [line.split(',') for line in text.strip().split('\n')]
        df = pd.DataFrame(data, columns=['id', 'text', 'label', 'metadata'])
        # 데이터 유형 최적화
        df = df.astype({'id': 'int32', 'text': 'string', 'label': 'string', 'metadata': 'string'})
        # Parquet 파일로 변환
        parquet_filename = 'text_to_parquet.parquet'
        df.to_parquet(parquet_filename, engine='pyarrow', compression='snappy')
        # Parquet 파일 내용 미리보기
        parquet_content = load_parquet(parquet_filename)
        return f"{parquet_filename} 파일이 성공적으로 변환되었습니다.", parquet_content, parquet_filename
    except Exception as e:
        return f"텍스트 변환 중 오류가 발생했습니다: {str(e)}", "", ""

# CSS 설정
css = """
footer {
    visibility: hidden;
}
#chatbot-container, #chatbot-data-upload {
    height: 700px;
    overflow-y: scroll;
}
#chatbot-container .message, #chatbot-data-upload .message {
    font-size: 14px;
}
/* 입력창 배경색 및 글자색 변경 */
textarea, input[type="text"] {
    background-color: #ffffff; /* 흰색 배경 */
    color: #000000; /* 검정색 글자 */
}
/* 파일 업로드 영역 높이 조절 */
#parquet-upload-area {
    max-height: 150px;
    overflow-y: auto;
}
/* 초기 설명 글씨 크기 조절 */
#initial-description {
    font-size: 14px;
}
"""


# Gradio Blocks 인터페이스 설정
with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
    gr.Markdown("# My RAG: LLM이 나만의 데이터로 학습한 콘텐츠 생성/답변", elem_id="initial-description")
    gr.Markdown(
        "### 1) 나만의 데이터를 입력 또는 CSV 업로드로 Parquet 데이터셋 자동 변환 2) Parquet 데이터셋을 업로드하면, LLM이 맞춤 학습 데이터로 활용하여 응답\n"
        "### Tip) '예제'를 통해 다양한 활용 방법을 체험하고 응용해 보세요, 데이터셋 업로드시 미리보기는 10건만 출력",
        elem_id="initial-description"
    )

    # 첫 번째 탭: 챗봇 데이터 업로드 (탭 이름 변경: "My 데이터셋+LLM")
    with gr.Tab("My 데이터셋+LLM"):
        gr.Markdown("### LLM과 대화하기")
        chatbot_data_upload = gr.Chatbot(label="챗봇", type="messages", elem_id="chatbot-data-upload")
        msg_data_upload = gr.Textbox(label="메시지 입력", placeholder="여기에 메시지를 입력하세요...")
        send_data_upload = gr.Button("전송")

        with gr.Accordion("시스템 프롬프트 및 옵션 설정", open=False):
            system_message = gr.Textbox(label="System Message", value="너는 AI 조언자 역할이다.")
            max_tokens = gr.Slider(minimum=1, maximum=8000, value=1000, label="Max Tokens")
            temperature = gr.Slider(minimum=0, maximum=1, value=0.7, label="Temperature")
            top_p = gr.Slider(minimum=0, maximum=1, value=0.9, label="Top P")

        parquet_data_state = gr.State()

        def handle_message_data_upload(
            message: str,
            history: List[Dict[str, str]],
            system_message: str,
            max_tokens: int,
            temperature: float,
            top_p: float,
            parquet_data: str
        ):
            history = history or []
            try:
                # 사용자의 메시지를 히스토리에 추가
                history.append({"role": "user", "content": message})
                # 응답 생성
                response_gen = respond(
                    message, history, system_message, max_tokens, temperature, top_p, parquet_data
                )
                partial_response = ""
                for partial in response_gen:
                    partial_response = partial
                    # 대화 내역 업데이트
                    display_history = history + [
                        {"role": "assistant", "content": partial_response}
                    ]
                    yield display_history, ""
                # 어시스턴트의 응답을 히스토리에 추가
                history.append({"role": "assistant", "content": partial_response})
            except Exception as e:
                response = f"추론 중 오류가 발생했습니다: {str(e)}"
                history.append({"role": "assistant", "content": response})
                yield history, ""

        send_data_upload.click(
            handle_message_data_upload,
            inputs=[
                msg_data_upload,
                chatbot_data_upload,
                system_message,
                max_tokens,
                temperature,
                top_p,
                parquet_data_state,  # parquet_data_state를 사용하여 업로드된 데이터를 전달
            ],
            outputs=[chatbot_data_upload, msg_data_upload],
            queue=True
        )

        # 예제 추가
        with gr.Accordion("예제", open=False):
            gr.Examples(
                examples=[
                    ["업로드된 데이터셋에 대해 요약 설명하라."],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 본 서비스를 SEO 최적화하여 블로그 포스트(개요, 배경 및 필요성, 기존 유사 제품/서비스와 비교하여 특장점, 활용처, 가치, 기대효과, 결론을 포함)로 4000 토큰 이상 작성하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 사용 방법과 차별점, 특징, 강점을 중심으로 4000 토큰 이상 유튜브 영상 스크립트 형태로 작성하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 제품 상세 페이지 형식의 내용을 4000 토큰 이상 자세히 설명하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, FAQ 20건을 상세하게 작성하라. 4000토큰 이상 사용하라."],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 특허 출원에 활용할 기술 및 비즈니스 모델 측면을 포함하여 특허 출원서 구성에 맞게 혁신적인 창의 발명 내용을 중심으로 4000 토큰 이상 작성하라."],
                ],
                inputs=msg_data_upload,
                label="예제 선택",
            )

        # Parquet 파일 업로드를 화면 하단으로 이동
        gr.Markdown("### Parquet 파일 업로드")
        with gr.Row():
            with gr.Column():
                parquet_upload = gr.File(
                    label="Parquet 파일 업로드", type="filepath", elem_id="parquet-upload-area"
                )
                parquet_upload_button = gr.Button("업로드")
                parquet_upload_status = gr.Textbox(label="업로드 상태", interactive=False)
                parquet_preview_chat = gr.Markdown(label="Parquet 파일 미리보기")

                def handle_parquet_upload(file_path: str):
                    message, parquet_content, parquet_json = upload_parquet(file_path)
                    if parquet_json:
                        return message, parquet_content, parquet_json
                    else:
                        return message, "", ""

                parquet_upload_button.click(
                    handle_parquet_upload,
                    inputs=parquet_upload,
                    outputs=[parquet_upload_status, parquet_preview_chat, parquet_data_state]
                )

    # 두 번째 탭: 데이터 변환 (탭 이름 변경: "CSV to My 데이터셋")
    with gr.Tab("CSV to My 데이터셋"):
        gr.Markdown("### CSV 파일 업로드 및 Parquet 변환")
        with gr.Row():
            with gr.Column():
                csv_file = gr.File(label="CSV 파일 업로드", type="filepath")
                upload_button = gr.Button("업로드 및 변환")
                upload_status = gr.Textbox(label="업로드 상태", interactive=False)
                parquet_preview = gr.Markdown(label="Parquet 파일 미리보기")
                download_button = gr.File(label="Parquet 파일 다운로드", interactive=False)

                def handle_csv_upload(file_path: str):
                    message, parquet_filename = upload_csv(file_path)
                    if parquet_filename:
                        parquet_content = load_parquet(parquet_filename)
                        return message, parquet_content, parquet_filename
                    else:
                        return message, "", None

                upload_button.click(
                    handle_csv_upload,
                    inputs=csv_file,
                    outputs=[upload_status, parquet_preview, download_button]
                )

    # 세 번째 탭: 텍스트 to csv to parquet 변환 (탭 이름 변경: "Text to My 데이터셋")
    with gr.Tab("Text to My 데이터셋"):
        gr.Markdown("### 텍스트를 입력하면 CSV로 변환 후 Parquet으로 자동 전환됩니다.")
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(
                    label="텍스트 입력 (각 행은 `id,text,label,metadata` 형식으로 입력)",
                    lines=10,
                    placeholder="예: 1,이순신,장군,거북선\n2,원균,장군,모함\n3,선조,왕,시기\n4,도요토미 히데요시,왕,침략"
                )
                convert_button = gr.Button("변환 및 다운로드")
                convert_status = gr.Textbox(label="변환 상태", interactive=False)
                parquet_preview_convert = gr.Markdown(label="Parquet 파일 미리보기")
                download_parquet_convert = gr.File(label="Parquet 파일 다운로드", interactive=False)

                def handle_text_to_parquet(text: str):
                    message, parquet_content, parquet_filename = text_to_parquet(text)
                    if parquet_filename:
                        return message, parquet_content, parquet_filename
                    else:
                        return message, "", None

                convert_button.click(
                    handle_text_to_parquet,
                    inputs=text_input,
                    outputs=[convert_status, parquet_preview_convert, download_parquet_convert]
                )

    gr.Markdown("### Arxivgpt@gmail.com", elem_id="initial-description")

if __name__ == "__main__":
    demo.launch()