import os
from huggingface_hub import HfApi
from tempfile import TemporaryDirectory
import importlib.util
import sys
import shutil
import subprocess

def load_private_code():
    token = os.environ.get("TOKEN", "")
    if token:
        try:
            subprocess.run(
                ["pip", "install", f"git+https://{token}@github.com/neoAI-inc/neo-llm-module.git@v1.3.25"],
                check=True
            )
        except subprocess.CalledProcessError as e:
            print("Error installing private module:", e)
    api = HfApi(token=os.environ["PRIVATE_REPOSITORY_HF_TOKEN"])
    
    tmp_dir = TemporaryDirectory()  # with文を使用しない


    try:
        # プライベートスペースからファイルをダウンロード
        api.snapshot_download(
            repo_id="neoai-inc-private/pj-andromeda",
            repo_type="space",
            local_dir=tmp_dir.name
        )
        
        # tmp_dirをPythonパスに追加
        sys.path.append(tmp_dir.name)
        
        # 現在のディレクトリを一時的に変更
        original_dir = os.getcwd()
        os.chdir(tmp_dir.name)
        
        # app.pyをインポート
        spec = importlib.util.spec_from_file_location(
            "private_app",
            os.path.join(tmp_dir.name, "app.py")
        )
        private_module = importlib.util.module_from_spec(spec)
        sys.modules["private_app"] = private_module
        
        try:
            spec.loader.exec_module(private_module)
        except ModuleNotFoundError as e:
            print(f"Error loading module: {e}")
            print(f"Current sys.path: {sys.path}")
            print(f"Contents of tmp_dir: {os.listdir(tmp_dir.name)}")
            raise
            
        # オリジナルのディレクトリに戻す
        os.chdir(original_dir)
        
        return private_module, tmp_dir  # tmp_dirも返す
    except Exception as e:
        raise e

# プライベートコードを読み込んでGradioアプリを取得
private_app, tmp_dir = load_private_code()  # tmp_dirを保持
app = private_app.app  # Gradioアプリのインスタンス

# Cleanup function for when the app stops
def cleanup():
    tmp_dir.cleanup()

# private側の認証情報を使用してアプリを起動
app.launch(
    auth=(private_app.USERNAME, private_app.PASSWORD),
    ssr_mode=False,
    root_path=tmp_dir.name  # 一時ディレクトリのパスを設定
)