import os
import torch
import onnx
import gradio as gr
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import HfApi, HfFolder, Repository

# モデルをONNXに変換する関数
def convert_to_onnx_custom_tensors(model_name, tensor_list, output_path="model.onnx"):
    # モデルとトークナイザをロード
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)

    # ユーザー指定のテンソル内容をパースして作成
    try:
        input_tensors = []
        for tensor_str in tensor_list:
            input_tensors.append(torch.tensor(eval(tensor_str)))
        
        # ダミーの例として、最初の2つのテンソルをinput_idsとattention_maskとして使用
        input_ids = input_tensors[0] if len(input_tensors) > 0 else torch.zeros((1, 1))
        attention_mask = input_tensors[1] if len(input_tensors) > 1 else torch.ones((1, 1))

    except Exception as e:
        return f"テンソル作成中にエラーが発生しました: {e}"

    # モデルをONNXにエクスポート
    try:
        torch.onnx.export(
            model,
            (input_ids, attention_mask),
            output_path,
            input_names=["input_ids", "attention_mask"],
            output_names=["output"],
            dynamic_axes={"input_ids": {0: "batch_size"}, "attention_mask": {0: "batch_size"}},
            opset_version=11
        )
        return f"モデルをONNX形式で{output_path}にエクスポートしました。"
    except Exception as e:
        return f"ONNX形式へのエクスポート中にエラーが発生しました: {e}"

# Hugging Face Hubにデプロイする関数
def deploy_to_huggingface(access_token, repo_name, model_path="model.onnx"):
    # Hugging Face Hubへの認証
    HfFolder.save_token(access_token)
    api = HfApi()

    # リポジトリを作成 (もし存在しない場合)
    try:
        if not api.model_info(repo_name, use_auth_token=access_token):
            api.create_repo(repo_name, exist_ok=True, token=access_token)

        # リポジトリにファイルをアップロード
        repo = Repository(local_dir=repo_name, clone_from=repo_name, use_auth_token=True)
        repo.git_pull()
        repo.lfs_track(["*.onnx"])
        os.makedirs(repo_name, exist_ok=True)
        model_output_path = os.path.join(repo_name, "model.onnx")
        os.rename(model_path, model_output_path)
        
        repo.git_add()
        repo.git_commit("Add ONNX model")
        repo.git_push()
        return f"{repo_name}にONNXモデルをデプロイしました。"
    except Exception as e:
        return f"デプロイ中にエラーが発生しました: {e}"

# AIの基本的な仕組みの解説
def ai_explanation():
    explanation = """
    AI（人工知能）は、人間の知的作業を模倣する技術です。
    特にディープラーニングはニューラルネットワークを用いて、大量のデータから特徴を学習し、分類や予測などを行うことができます。
    モデルはトレーニングフェーズでデータを用いて学習し、その後の推論フェーズで新しいデータに対して応答を生成します。
    ONNXは異なるフレームワーク間で互換性を持たせるための形式で、PyTorchやTensorFlowのモデルを統一的に使うことができます。
    """
    return explanation

# Gradioインターフェースの作成
def main():
    with gr.Blocks() as demo:
        gr.Markdown("# ONNXモデル変換・デプロイツールとAIの解説")
        
        # AIの基本説明セクション
        with gr.Tab("AIの基本説明"):
            explanation = gr.Markdown(ai_explanation())

        # ONNX変換セクション
        with gr.Tab("ONNXモデルの変換"):
            model_name = gr.Textbox(label="モデル名", placeholder="例: bert-base-uncased")
            tensor_list = gr.Dataframe(headers=["テンソルの内容"], datatype="str", row_count=3, col_count=1, placeholder="例: [[101, 2057, 2024, 102]]")
            convert_btn = gr.Button("ONNXに変換")
            output = gr.Textbox(label="出力メッセージ")
            
            convert_btn.click(
                convert_to_onnx_custom_tensors,
                inputs=[model_name, tensor_list],
                outputs=output
            )
        
        # デプロイセクション
        with gr.Tab("Hugging Faceにデプロイ"):
            access_token = gr.Textbox(label="Hugging Faceアクセストークン", type="password")
            repo_name = gr.Textbox(label="リポジトリ名")
            deploy_btn = gr.Button("デプロイ")
            deploy_output = gr.Textbox(label="デプロイ出力")

            deploy_btn.click(
                deploy_to_huggingface,
                inputs=[access_token, repo_name],
                outputs=deploy_output
            )

    demo.launch()

if __name__ == "__main__":
    main()