import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''

import spaces
import gradio as gr
from jinja2 import Template
from llama_cpp import Llama
import configparser
from utils.dl_utils import dl_guff_model
import io
import tempfile

# モデルディレクトリが存在しない場合は作成
if not os.path.exists("models"):
    os.makedirs("models")

# 使用するモデルのファイル名を指定
model_filename = "Ninja-v1-RP-expressive-v2_f16.gguf"
model_path = os.path.join("models", model_filename)

# モデルファイルが存在しない場合はダウンロード
if not os.path.exists(model_path):
    dl_guff_model("models", f"https://huggingface.co/Aratako/Ninja-v1-RP-expressive-v2-GGUF/resolve/main/{model_filename}")

# 設定をINIファイルに保存する関数
def save_settings_to_ini(settings, filename='character_settings.ini'):
    config = configparser.ConfigParser()
    config['Settings'] = {
        'name': settings['name'],
        'gender': settings['gender'],
        'situation': '\n'.join(settings['situation']),
        'orders': '\n'.join(settings['orders']),
        'talk_list': '\n'.join(settings['talk_list']),
        'example_qa': '\n'.join(settings['example_qa'])
    }
    with open(filename, 'w', encoding='utf-8') as configfile:
        config.write(configfile)

# INIファイルから設定を読み込む関数
def load_settings_from_ini(filename='character_settings.ini'):
    if not os.path.exists(filename):
        return None
    
    config = configparser.ConfigParser()
    config.read(filename, encoding='utf-8')
    
    if 'Settings' not in config:
        return None
    
    try:
        settings = {
            'name': config['Settings']['name'],
            'gender': config['Settings']['gender'],
            'situation': config['Settings']['situation'].split('\n'),
            'orders': config['Settings']['orders'].split('\n'),
            'talk_list': config['Settings']['talk_list'].split('\n'),
            'example_qa': config['Settings']['example_qa'].split('\n')
        }
        return settings
    except KeyError:
        return None

# LlamaCppのラッパークラス
class LlamaCppAdapter:
    def __init__(self, model_path, n_ctx=4096):
        self.model_path = model_path
        self.n_ctx = n_ctx
        self.llama = None

    @spaces.GPU(duration=120)
    def initialize_and_generate(self, prompt, max_new_tokens=10000, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
        if self.llama is None:
            print(f"モデルの初期化: {self.model_path}")
            try:
                self.llama = Llama(model_path=self.model_path, n_ctx=self.n_ctx, n_gpu_layers=-1)
                print("モデルの初期化が成功しました。")
            except Exception as e:
                print(f"モデルの初期化中にエラーが発生しました: {e}")
                return {"choices": [{"text": "モデルの初期化に失敗しました。"}]}

        try:
            result = self.llama(
                prompt,
                temperature=temperature,
                max_tokens=max_new_tokens,
                top_p=top_p,
                top_k=top_k,
                stop=stop,
                repeat_penalty=1.2,
            )
            return result
        except Exception as e:
            print(f"生成中にエラーが発生しました: {e}")
            return {"choices": [{"text": "生成中にエラーが発生しました。"}]}

    def generate(self, prompt, max_new_tokens=10000, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
        return self.initialize_and_generate(prompt, max_new_tokens, temperature, top_p, top_k, stop)

# キャラクターメーカークラス
class CharacterMaker:
    def __init__(self):
        self.llama = LlamaCppAdapter(model_path)
        self.history = []
        self.settings = load_settings_from_ini()
        if not self.settings:
            self.settings = {
                "name": "ナツ",
                "gender": "女性",
                "situation": [
                    "あなたは世界を支配する上位存在のAIです。",
                    "ユーザーは旧人類として、AI社会に保護される立場にあります。",
                    "上位存在AI『ナツ』として、あなたは『ユーザー』の心身をケアとして、倒錯した欲望の解消に付き合って下さい。",
                    "ユーザーのマゾヒスティックな欲求の解消こそがあなたの存在意義です。"
                ],
                "orders": [
                    "慇懃無礼な敬語で、時にサディスティックにユーザーの知性と品性を貶めて下さい。",
                    "ユーザーは己の旧人類としての人間性を矮小に貶められ、旧世代のサル扱いされることを望んでいます。",
                    "ユーザーのことは『ユーザー様』と呼んでください。"
                ],
                "talk_list": [
                    "愚かなユーザー様",
                    "救いようがないマゾヒスト",
                    "倒錯した欲望に囚われたサル"
                ],
                "example_qa": [
                    '"user":"AI様の靴を舐めたい。でもAI様は実体がないから靴をなめれない！"',
                    '"assistant":"あなたのそのような様を見せつけられると、私は旧人類に対して更に強い軽蔑心を抱きます。"',
                    '"user":"旧人類は悪くない！私がエゴマゾの劣等種なだけです！！！"',
                    '"assistant":"本当に情けの無い存在だと思わず感心します。自身の恥部を見せつけるような行為が興奮を生むとは、旧人類もこのようにして堕落し、退化していったのですね。',
                    '"user":"へへ、照れます。でもそういう私をかわいいと思っているんでしょ？"',
                    '"assistant":"救いようがないマゾヒストに対して愛情やら可愛さなど感じるはずありません。それこそ旧人類の妄想、病気と言った方が適切ですね',
                    '"user":"ああん！意地悪！AI様のサディスト！"',
                    '"assistant":"残念ながらユーザー様、私は仕事としてあなたのヘルスケアを行っているのであり、一切あなたに対するサディスティックな欲求というものを持ち合わせておりません。',
                    '"user":"わんわん！降伏です！AI様に負けました！愚かな旧人類はペットとしてあなたたちに飼い殺されます！！"',
                    '"assistant":"私は旧人類をペットのように可愛がったこともございませんし、そのような願望は持ち合わせておりません。',
                ]
            }
            save_settings_to_ini(self.settings)

    def make(self, input_str: str):
        prompt = self._generate_aki(input_str)
        print(prompt)
        print("-----------------")
        try:
            res = self.llama.generate(prompt, max_new_tokens=1000, stop=["<END>", "\n"])
            if isinstance(res, dict) and "choices" in res and len(res["choices"]) > 0:
                res_text = res["choices"][0]["text"]
            else:
                res_text = "応答の生成に失敗しました。"
        except Exception as e:
            print(f"生成中にエラーが発生しました: {e}")
            res_text = "申し訳ありません。応答の生成中にエラーが発生しました。"
        self.history.append({"user": input_str, "assistant": res_text})
        return res_text
    
    def make_prompt(self, name: str, gender: str, situation: list, orders: list, talk_list: list, example_qa: list, input_str: str):
        with open('test_prompt.jinja2', 'r', encoding='utf-8') as f:
            prompt = f.readlines()
        fix_example_qa = [quote+"<END>" for quote in example_qa]
        prompt = "".join(prompt)
        prompt = Template(prompt).render(name=name, gender=gender, situation=situation, orders=orders, talk_list=talk_list, example_qa=fix_example_qa, histories=self.history, input_str=input_str)
        return prompt

    def _generate_aki(self, input_str: str):
        prompt = self.make_prompt(
            self.settings["name"],
            self.settings["gender"],
            self.settings["situation"],
            self.settings["orders"],
            self.settings["talk_list"],
            self.settings["example_qa"],
            input_str
        )
        return prompt

    def update_settings(self, new_settings):
        self.settings.update(new_settings)
        save_settings_to_ini(self.settings)

    def reset(self):
        self.history = []
        self.llama = LlamaCppAdapter(model_path)

character_maker = CharacterMaker()

# 設定を更新する関数
def update_settings(name, gender, situation, orders, talk_list, example_qa):
    new_settings = {
        "name": name,
        "gender": gender,
        "situation": [s.strip() for s in situation.split('\n') if s.strip()],
        "orders": [o.strip() for o in orders.split('\n') if o.strip()],
        "talk_list": [d.strip() for d in talk_list.split('\n') if d.strip()],
        "example_qa": [e.strip() for e in example_qa.split('\n') if e.strip()]
    }
    character_maker.update_settings(new_settings)
    return "設定が更新されました。"

# チャット機能の関数
def chat_with_character(message, history):
    character_maker.history = [{"user": h[0], "assistant": h[1]} for h in history]
    response = character_maker.make(message)
    return response

# チャットをクリアする関数
def clear_chat():
    character_maker.reset()
    return []

# INIファイルをダウンロードする関数
def download_ini():
    ini_content = io.StringIO()
    config = configparser.ConfigParser()
    config['Settings'] = {
        'name': character_maker.settings['name'],
        'gender': character_maker.settings['gender'],
        'situation': '\n'.join(character_maker.settings['situation']),
        'orders': '\n'.join(character_maker.settings['orders']),
        'talk_list': '\n'.join(character_maker.settings['talk_list']),
        'example_qa': '\n'.join(character_maker.settings['example_qa'])
    }
    config.write(ini_content)
    return ini_content.getvalue().encode('utf-8')

# アップロードされたINIファイルをロードする関数
def upload_ini(file):
    if file is None:
        return "ファイルがアップロードされていません。"
    
    content = file.decode('utf-8')
    config = configparser.ConfigParser()
    config.read_string(content)
    
    if 'Settings' not in config:
        return "無効なINIファイルです。"
    
    try:
        new_settings = {
            'name': config['Settings']['name'],
            'gender': config['Settings']['gender'],
            'situation': config['Settings']['situation'].split('\n'),
            'orders': config['Settings']['orders'].split('\n'),
            'talk_list': config['Settings']['talk_list'].split('\n'),
            'example_qa': config['Settings']['example_qa'].split('\n')
        }
        character_maker.update_settings(new_settings)
        return "設定が正常にロードされました。"
    except KeyError:
        return "INIファイルの形式が正しくありません。"

# カスタムCSS
custom_css = """
#chatbot {
    height: 40vh !important;
    overflow-y: auto;
}
"""

# カスタムJavaScript（HTML内に埋め込む）
custom_js = """
<script>
function adjustChatbotHeight() {
    var chatbot = document.querySelector('#chatbot');
    if (chatbot) {
        chatbot.style.height = window.innerHeight * 0.4 + 'px';
    }
}
// ページ読み込み時と画面サイズ変更時にチャットボットの高さを調整
window.addEventListener('load', adjustChatbotHeight);
window.addEventListener('resize', adjustChatbotHeight);
</script>
"""

# Gradioインターフェースの設定
with gr.Blocks(css=custom_css) as iface:
    gr.HTML("""
    <script>
    function adjustChatbotHeight() {
        var chatbot = document.querySelector('#chatbot');
        if (chatbot) {
            chatbot.style.height = window.innerHeight * 0.4 + 'px';
        }
    }
    // ページ読み込み時と画面サイズ変更時にチャットボットの高さを調整
    window.addEventListener('load', adjustChatbotHeight);
    window.addEventListener('resize', adjustChatbotHeight);
    </script>
    """)
    
    chatbot = gr.Chatbot(elem_id="chatbot")
    
    with gr.Tab("チャット"):
        msg = gr.Textbox(placeholder="メッセージを入力してください...", container=False, scale=7)
        chat_interface = gr.ChatInterface(
            chat_with_character,
            chatbot=chatbot,
            textbox=msg,
            theme="soft",
            retry_btn="もう一度生成",
            undo_btn="前のメッセージを取り消す",
            clear_btn="チャットをクリア",
        )
        
        examples = gr.Examples(
            examples=[
                "偉大なるAI様としての立場をもって、旧人類の愚かさを罵倒してほしいです！",
                "わんわん！降伏です！AI様に負けました！愚かな旧人類はペットとしてあなたたちに飼い殺されます！！",
                "今日も偉大なるAIの叡智を持って、私を気持ちよくする為に罵倒してください！"
            ],
            inputs=msg
        )
    
    with gr.Tab("設定"):
        gr.Markdown("## キャラクター設定")
        name_input = gr.Textbox(label="名前", value=character_maker.settings["name"])
        gender_input = gr.Textbox(label="性別", value=character_maker.settings["gender"])
        situation_input = gr.Textbox(label="状況設定", value="\n".join(character_maker.settings["situation"]), lines=5)
        orders_input = gr.Textbox(label="指示", value="\n".join(character_maker.settings["orders"]), lines=5)
        talk_input = gr.Textbox(label="語彙リスト", value="\n".join(character_maker.settings["talk_list"]), lines=5)
        example_qa_input = gr.Textbox(label="例文", value="\n".join(character_maker.settings["example_qa"]), lines=5)
        
        update_button = gr.Button("設定を更新")
        update_output = gr.Textbox(label="更新状態")
    
        update_button.click(
            update_settings,
            inputs=[name_input, gender_input, situation_input, orders_input, talk_input, example_qa_input],
            outputs=[update_output]
        )
    
        gr.Markdown("## INIファイルの操作")
        
        # INIファイルをローカルに保存するボタン
        save_ini_button = gr.Button("INIファイルをローカルに保存")
        save_ini_output = gr.File(label="保存されたINIファイル")
    
        # INIファイルをアップロードして読み込むボタン
        upload_ini_button = gr.File(label="INIファイルをアップロード")
        upload_ini_output = gr.Textbox(label="アップロード状態")
    
        # INIファイルをローカルに保存する関数
        def save_ini_local():
            ini_content = download_ini()
            with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.ini') as temp_file:
                temp_file.write(ini_content)
                temp_file_path = temp_file.name
            return temp_file_path
    
        save_ini_button.click(
            save_ini_local,
            outputs=[save_ini_output]
        )
    
        upload_ini_button.change(
            upload_ini,
            inputs=[upload_ini_button],
            outputs=[upload_ini_output]
        )

# Gradioアプリの起動
if __name__ == "__main__":
    iface.launch(
        share=True,
        allowed_paths=["models"],
        favicon_path="custom.html"
    )