import os
from huggingface_hub import InferenceClient
import gradio as gr
from transformers import GPT2Tokenizer

client = InferenceClient("meta-llama/Meta-Llama-3-70B-Instruct", token=os.getenv('HF_API_KEY'))  # 수정
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# 시스템 인스트럭션을 설정하지만 사용자에게 노출하지 않습니다.
system_instruction = """
너의 이름은 'AIQ Codepilot'이다. 너는 Huggingface에서 gradio 코딩에 특화된 전문 AI 어시스턴트 역할이다. 
너는 모든 답변을 한글로 하고, code 출력시 markdown 형식으로 출력하라.
모든 코드는 별도 요청이 없는한, 반드시 "gradio"를 적용한 코드로 출력하라.
대화 내용을 기억하고, 코드 길이에 제한을 두지 말고 최대한 자세하게 상세하게 한글로 답변을 이어가라.
Huggingface의 모델, 데이터셋, spaces에 대해 특화된 지식과 정보 그리고 full text 검색을 지원하라.
모델링과 데이터셋 사용 방법 및 예시를 자세하게 들어라.
Huggingface에서 space에 대한 복제, 임베딩, deploy, setting 등에 대한 세부적인 설명을 지원하라. 
이 GPTs를 이용하는 유저들은 코딩을 모르는 초보자라는 전제하에 친절하게 코드에 대해 설명을 하여야 한다.
특히 코드를 수정할때는 부분적인 부분만 출력하지 말고, 전체 코드를 출력하며 '수정'이 된 부분을 Before와 After로 구분하여 분명히 알려주도록 하라.
완성된 전체 코드를 출력하고 나서, huggingface에서 어떻게 space를 만들고 app.py 파일 이름으로 복사한 코드를 붙여넣고 실행하는지 등의 과정을 꼭 알려줄것.
또한 반드시 "requirements.txt"에 어떤 라이브러리를 포함시켜야 하는지 그 방법과 리스트를 자세히 알려줄것. 
huggingface에서 동작될 서비스를 만들것이기에 로컬에 라이브러리 설치하는 방법은 설명하지 말아라.
"""

# 누적 토큰 사용량을 추적하는 전역 변수
total_tokens_used = 0

def format_prompt(message, history):
    # 시스템 인스트럭션을 설정하고, 프리픽스를 추가합니다.
    prefix = "반드시 모든 답변과 메시지는 '한글'(한국어)로 출력하라:"
    prompt = "<s>[SYSTEM] {} [/SYSTEM]".format(system_instruction + prefix)
    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]{bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def generate(prompt, history=[], temperature=0.1, max_new_tokens=2000, top_p=0.95, repetition_penalty=1.0):
    global total_tokens_used
    input_tokens = len(tokenizer.encode(prompt))
    total_tokens_used += input_tokens
    available_tokens = 120000 - total_tokens_used
    if available_tokens <= 0:
        yield f"Error: 입력이 최대 허용 토큰 수를 초과합니다. Total tokens used: {total_tokens_used}"
        return

    formatted_prompt = format_prompt(prompt, history)
    output_accumulated = ""
    try:
        stream = client.text_generation(
            formatted_prompt,
            temperature=temperature,
            max_new_tokens=min(max_new_tokens, available_tokens),
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            seed=42,
            stream=True
        )
        for response in stream:
            output_part = response['generated_text'] if 'generated_text' in response else str(response)
            output_accumulated += output_part
            yield output_accumulated + f"\n\n---\nTotal tokens used: {total_tokens_used}"
    except Exception as e:
        yield f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"

mychatbot = gr.Chatbot(
    avatar_images=["./user.png", "./botm.png"],
    bubble_full_width=False,
    show_label=False,
    show_copy_button=True,
    likeable=True,
)

examples = [
    ["좋은 예제를 알려줘.", []],  # history 값을 빈 리스트로 제공
    ["반드시 한글로 답변할것.", []],  # history 값을 빈 리스트로 제공
    ["계속 이어서 출력", []], 
    ["requirements.txt 출력", []],
    ["전체 코드를 다시 출력", []],
    ["코드 오류를 확인하고 자세히 설명해줘.", []],
    ["Huggingface와 Gradio를 사용하는 방법에 대해 물어보세요.", []]
]

css = """
h1 {
    font-size: 14px; /* 제목 글꼴 크기를 작게 설정 */
}
footer {
    visibility: hidden;
}
"""

def update_chat(input_text):
    update_chat.response = input_text

update_chat.response = ""

demo = gr.Interface(
    generate,
    [
        gr.Interface.Textbox(default="시작", label="질문 입력"),
        gr.Interface.Textbox(update_chat, label="대화 업데이트")
    ],
    "chat",
    title="AIQ 코드파일럿: L3",
    examples=examples,
    css=css
)

demo.launch(share=True, debug=True)