import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import html

# モデルのダウンロードと初期化
def download_model():
    model_name = "bartowski/DeepSeek-Coder-V2-Lite-Instruct-GGUF"
    model_file = "DeepSeek-Coder-V2-Lite-Instruct-Q6_K.gguf"
    try:
        return hf_hub_download(repo_id=model_name, filename=model_file)
    except Exception as e:
        print(f"Model download failed: {e}")
        return None

def initialize_model():
    try:
        model_path = download_model()
        if not model_path:
            return None
        return Llama(
            model_path=model_path,
            n_ctx=4096,
            n_threads=4,
            n_gpu_layers=-1
        )
    except Exception as e:
        print(f"Error initializing model: {e}")
        return None

llm = initialize_model()

system_prompt = (
    "You are a helpful AI coding assistant. Your mission is to help people with programming "
    "and technical questions, providing clear and concise answers."
)

# チャットボットの処理関数
def respond(message, chat_history):
    if not message or not llm:
        return chat_history
    
    # チャット履歴の構築（システムプロンプト + 過去の会話 + 新しいメッセージ）
    messages = [{"role": "system", "content": system_prompt}]
    
    for user_msg, bot_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})
    
    # AIからのレスポンス生成
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=1000,
        stop=["User:"],
        stream=False
    )
    
    ai_response = response['choices'][0]['message']['content']
    
    # チャット履歴に追加
    chat_history.append((message, ai_response))
    
    return chat_history

# Gradioインターフェースの作成
with gr.Blocks() as demo:
    gr.Markdown("# AI Coding Assistant")
    
    chatbot = gr.Chatbot(height=500)
    msg = gr.Textbox(label="Your Message")
    clear = gr.Button("Clear")
    
    msg.submit(respond, [msg, chatbot], chatbot)
    clear.click(lambda: None, None, chatbot, queue=False)

# アプリの起動
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)