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

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

def load_code(filename):
    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):
    try:
        df = pd.read_parquet(filename, engine='pyarrow')
        # 데이터프레임의 첫 몇 행을 문자열로 변환 (Markdown 표 형식)
        return df.head(10).to_markdown(index=False)
    except FileNotFoundError:
        return f"{filename} 파일을 찾을 수 없습니다."
    except Exception as e:
        return f"파일을 읽는 중 오류가 발생했습니다: {str(e)}"

# 코드 파일 로드
fashion_code = load_code('fashion.cod')
uhdimage_code = load_code('uhdimage.cod')
MixGEN_code = load_code('mgen.cod')

# 초기 Parquet 파일 로드 (기존 test.parquet)
test_parquet_content = load_parquet('test.parquet')

def respond(
    message,
    history: List[Tuple[str, str]],
    system_message="",  # 기본값 추가
    max_tokens=4000,    # 기본값 변경
    temperature=0.7,    # 기본값 유지
    top_p=0.9,           # 기본값 유지
    ):
    # 시스템 프롬프트 설정
    system_prefix = """반드시 한글로 답변할것. 너는 주어진 소스코드를 기반으로 "서비스 사용 설명 및 안내, Q&A를 하는 역할이다". 아주 친절하고 자세하게 4000토큰 이상 Markdown 형식으로 작성하라. 너는 코드를 기반으로 사용 설명 및 질의 응답을 진행하며, 이용자에게 도움을 주어야 한다. 이용자가 궁금해 할 만한 내용에 친절하게 알려주도록 하라. 코드 전체 내용에 대해서는 보안을 유지하고, 키 값 및 엔드포인트와 구체적인 모델은 공개하지 마라."""
    
    # 명령어 처리
    if message.lower() == "패션 코드 실행":
        system_message += f"\n\n패션 코드 내용:\n```python\n{fashion_code}\n```"
        message = "패션 가상피팅에 대한 내용을 학습하였고, 설명할 준비가 되어있다고 알리고 서비스 URL(https://aiqcamp-fash.hf.space)을 통해 테스트 해보라고 출력하라."
    elif message.lower() == "uhd 이미지 코드 실행":
        system_message += f"\n\nUHD 이미지 코드 내용:\n```python\n{uhdimage_code}\n```"
        message = "UHD 이미지 생성에 대한 내용을 학습하였고, 설명할 준비가 되어있다고 알리고 서비스 URL(https://openfree-ultpixgen.hf.space)을 통해 테스트 해보라고 출력하라."
    elif message.lower() == "mixgen 코드 실행":
        system_message += f"\n\nMixGEN 코드 내용:\n```python\n{MixGEN_code}\n```"
        message = "MixGEN3 이미지 생성에 대한 내용을 학습하였고, 설명할 준비가 되어있다고 알리고 서비스 URL(https://openfree-mixgen3.hf.space)을 통해 테스트 해보라고 출력하라."
    elif message.lower() == "test.parquet 실행":
        # 현재 Parquet 내용은 상태 변수에서 가져와야 함
        # 여기서는 예시로 'parquet_content'가 history에 있다고 가정
        parquet_content = ""
        for item in history:
            if isinstance(item, tuple) and item[0].startswith("Parquet"):
                parquet_content = item[1]
                break
        system_message += f"\n\ntest.parquet 파일 내용:\n```markdown\n{parquet_content}\n```"
        message = "test.parquet 파일에 대한 내용을 학습하였고, 관련 설명 및 Q&A를 진행할 준비가 되어있다. 궁금한 점이 있으면 물어보라."
    elif message.lower() == "csv 업로드":
        message = "CSV 파일을 업로드하려면 두 번째 탭을 사용하세요."

    # 시스템 메시지와 사용자 메시지 결합
    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]
    for val in history:
        if isinstance(val, tuple) and val[0]:
            messages.append({"role": "user", "content": val[0]})
        if isinstance(val, tuple) and val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    response = ""
    try:
        for msg in hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = msg.choices[0].delta.get('content', None)
            if token:
                response += token
                yield response
    except Exception as e:
        yield f"추론 중 오류가 발생했습니다: {str(e)}"

def upload_csv(file):
    try:
        # CSV 파일 읽기 (구분자 콤마)
        df = pd.read_csv(file, sep=',')

        # CSV 파일의 컬럼 확인
        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)}", None

        # 데이터 클렌징
        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))[0] + '.parquet'
        df.to_parquet(parquet_filename, engine='pyarrow', compression='snappy')

        # Parquet 파일 로드
        parquet_content = load_parquet(parquet_filename)

        return f"{parquet_filename} 파일이 성공적으로 업로드되고 변환되었습니다.", parquet_filename
    except Exception as e:
        return f"CSV 파일 업로드 및 변환 중 오류가 발생했습니다: {str(e)}", None

def upload_parquet(file):
    try:
        # Parquet 파일 읽기
        df = pd.read_parquet(file, engine='pyarrow')

        # 데이터프레임을 Markdown으로 변환
        parquet_content = df.to_markdown(index=False)

        return "Parquet 파일이 성공적으로 업로드되었습니다.", parquet_content, df.to_json()  # JSON으로 데이터 저장
    except Exception as e:
        return f"Parquet 파일 업로드 중 오류가 발생했습니다: {str(e)}", None, None

css = """
footer {
    visibility: hidden;
}
"""

# Gradio Blocks 인터페이스 설정
with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
    gr.Markdown("# LLM 서비스 인터페이스")

    with gr.Tab("챗봇"):
        gr.Markdown("### LLM과 대화하기")
        chatbot = gr.Chatbot(label="챗봇")
        msg = gr.Textbox(label="메시지 입력")
        send = 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=4000, 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")

        # 챗봇 메시지 처리 함수
        def handle_message(message, history, system_message, max_tokens, temperature, top_p):
            # Update history with user message
            history = history or []
            history.append((message, None))
            # Generate response
            response = ""
            try:
                for token in respond(message, history, system_message, max_tokens, temperature, top_p):
                    response = token
                # Update history with assistant response
                history[-1] = (message, response)
            except Exception as e:
                history[-1] = (message, f"추론 중 오류가 발생했습니다: {str(e)}")
            return history, ""

        send.click(
            handle_message,
            inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p],
            outputs=[chatbot, msg]
        )

    with gr.Tab("데이터 변환"):
        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 파일 다운로드", type="binary", interactive=False)

                # 업로드 버튼 클릭 시 실행할 함수
                def handle_csv_upload(file):
                    message, parquet_filename = upload_csv(file)
                    if parquet_filename:
                        # 파일을 다운로드할 수 있도록 경로 설정
                        with open(parquet_filename, "rb") as f:
                            data = f.read()
                        return message, load_parquet(parquet_filename), data
                    else:
                        return message, "", None

                upload_button.click(
                    handle_csv_upload,
                    inputs=csv_file,
                    outputs=[upload_status, parquet_preview, download_button]
                )
        
        gr.Markdown("### 기존 Parquet 파일")
        gr.Markdown(f"**test.parquet 파일 내용:**\n```markdown\n{test_parquet_content}\n```")

    with gr.Tab("챗봇 데이터 업로드"):
        gr.Markdown("### Parquet 파일 업로드 및 질문하기")
        with gr.Row():
            with gr.Column():
                parquet_upload = gr.File(label="Parquet 파일 업로드", type="filepath")
                parquet_upload_button = gr.Button("업로드")
                parquet_upload_status = gr.Textbox(label="업로드 상태", interactive=False)
                parquet_preview_chat = gr.Markdown(label="Parquet 파일 미리보기")
                # 상태를 저장할 Hidden State
                parquet_data_state = gr.State()

                def handle_parquet_upload(file):
                    message, parquet_content, parquet_json = upload_parquet(file)
                    if parquet_json:
                        return message, parquet_preview_chat.update(value=parquet_content), parquet_data_state.update(value=parquet_json)
                    else:
                        return message, gr.Markdown.update(value=""), parquet_data_state.update(value=None)

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

        gr.Markdown("### LLM과 대화하기")
        chatbot_data_upload = gr.Chatbot(label="챗봇 데이터 업로드")
        msg_data_upload = gr.Textbox(label="메시지 입력")
        send_data_upload = gr.Button("전송")

        # 챗봇 메시지 처리 함수 (데이터 업로드 버전)
        def handle_message_data_upload(message, history, system_message, max_tokens, temperature, top_p, parquet_data):
            # Update history with user message
            history = history or []
            history.append((message, None))
            # Generate response
            response = ""
            try:
                for token in respond(message, history, system_message, max_tokens, temperature, top_p):
                    response = token
                # Update history with assistant response
                history[-1] = (message, response)
            except Exception as e:
                history[-1] = (message, f"추론 중 오류가 발생했습니다: {str(e)}")
            return 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],
            outputs=[chatbot_data_upload, msg_data_upload]
        )

    gr.Markdown("## 주의 사항")
    gr.Markdown("""
    - **CSV 업로드**: CSV 파일을 업로드하면 자동으로 Parquet 파일로 변환됩니다. CSV 파일은 반드시 **콤마(`,`)**로 구분되어야 합니다.
    - **Parquet 미리보기**: 업로드된 Parquet 파일의 첫 10개 행이 미리보기로 표시됩니다.
    - **LLM과의 대화**: 변환된 Parquet 파일 내용을 기반으로 LLM이 응답을 생성합니다.
    - **Parquet 다운로드**: 변환된 Parquet 파일을 다운로드하려면 변환된 파일 옆의 다운로드 링크를 클릭하세요.
    - **챗봇 데이터 업로드**: 챗봇 데이터 업로드 탭에서 Parquet 파일을 업로드하면 해당 데이터를 기반으로 질문과 답변을 진행할 수 있습니다.
    """)

    gr.Markdown("### Gradio 인터페이스를 사용하여 LLM 모델과 상호작용하세요!")

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