# -*- coding: utf-8 -*-

import gradio as gr
from huggingface_hub import InferenceClient
from gradio_client import Client
import os
import requests
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

# API 설정
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN"))
IMAGE_API_URL = "http://211.233.58.201:7896"

def generate_image(prompt: str) -> tuple:
    """이미지 생성 함수"""
    try:
        client = Client(IMAGE_API_URL)
        # 프롬프트 앞에 "fantasy style," 추가
        enhanced_prompt = f"fantasy style, {prompt}"
        result = client.predict(
            prompt=enhanced_prompt,
            width=768,
            height=768,
            guidance=7.5,
            inference_steps=30,
            seed=3,
            do_img2img=False,
            init_image=None,
            image2image_strength=0.8,
            resize_img=True,
            api_name="/generate_image"
        )
        return result[0], result[1]
    except Exception as e:
        logging.error(f"Image generation failed: {str(e)}")
        return None, f"Error: {str(e)}"

def respond(
    message,
    history: list[tuple[str, str]],
    system_message="",  
    max_tokens=7860,  
    temperature=0.8, 
    top_p=0.9,  
):
    system_prefix = """
    [시스템 프롬프트 내용...]
    """

    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    response = ""
    try:
        for message in hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content
            if token is not None:
                response += token.strip("")
            yield response, None  # 이미지를 위한 None 추가

        # 텍스트 생성이 완료된 후 이미지 생성
        image, seed = generate_image(response[:200])  # 처음 200자를 이미지 프롬프트로 사용
        yield response, image

    except Exception as e:
        yield f"Error: {str(e)}", None

# Gradio 인터페이스 설정
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as interface:
    gr.Markdown("# Fantasy Novel AI Generation")
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="Enter your message")
            system_msg = gr.Textbox(label="System Message", value="Write(output) in 한국어.")
            
            with gr.Row():
                max_tokens = gr.Slider(minimum=1, maximum=8000, value=7000, 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")
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image")

    examples = gr.Examples(
        examples=[
            ["판타지 소설의 흥미로운 소재 10가지를 제시하라"],
            ["계속 이어서 작성하라"],
            ["Translate into English"],
            ["마법 시스템에 대해 더 자세히 설명하라"],
            ["전투 장면을 더 극적으로 묘사하라"],
            ["새로운 판타지 종족을 추가하라"],
            ["고대 예언에 대해 더 자세히 설명하라"],
            ["주인공의 내면 묘사를 추가하라"],
        ],
        inputs=msg
    )

    msg.submit(
        respond,
        [msg, chatbot, system_msg, max_tokens, temperature, top_p],
        [chatbot, image_output]
    )

# 애플리케이션 실행
if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )