# -*- coding: utf-8 -*-
"""QVP2_20250401_01.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Nef_3Xp9k0z8jV-ezGjHjTQDKn8mP7TJ

# QuickViper2
Conda実行環境権限不足に対応するため、

1. 仮想環境の圧縮書庫を作成
2. 1.でできた圧縮書庫を指定フォルダに展開

を行う。



```
2025/03/22 0.1.0 完成
2025/03/22 0.1.1 バグ修正
2025/03/23 0.1.3 DebugHelperとの依存関係修復のため
2025/03/23 0.1.6 バグ修正
2025/03/25 0.1.7 バグ修正
2025/03/25 0.1.8 各種修正と改良。UniqueEnvVar を導入
2025/03/26 0.2.0 ブラッシュアップ。EasyVen, DebugHelper.installer導入
2025/03/26 0.2.1 conda_unpack修正。仮想環境をアクティブにできるようになる。
2025/03/27 0.2.2 compress_venv バグ修正
2025/03/31 1.0.0 第一次改装 CondaInitializer追加 他
2025/03/31 1.0.1 ミス
2025/04/01 1.1.0 pack, extract 改良など
2025/04/01 1.1.1 extract修正
```

# QuickViper2

## 1. condacolab 起動
### CondaInitializer
- **概要:** Google Drive接続と condacolabのインストールを行う
- **使用例:** コード冒頭に設置。condacolabと周辺プログラムが起動する。
"""

# @title a. CondaInitializer 定義 & 実行
from google.colab import drive
drive.mount('/content/drive')
# condacolab 起動
get_ipython().system( "pip install -q condacolab" )
import condacolab
get_ipython().system( "pip install an_DebugHelper" )
from an_debughelper import DebugHelper


class CondaInitializer:
    def __init__( self ):
        """
        condaの初期設定を行います。
        """
        self.debug = DebugHelper( instance_name = "CondaInitializer")
        self.debug.enable_log_to_file_stdout()
        self.debug.enable_log_to_file_stderr()
        self.debug.enable_timestamp()

        # GDrive接続

        # Colabでは "source" を使えないため、bash -c で conda 初期化スクリプトを経由する
        self.debug.installer_sync( command = [ "bash", "-c", "source", "/usr/local/etc/profile.d/conda.sh" ], shell = True )

        condacolab.install()
        self.debug.installer_sync( "conda install -q mamba -n base -c conda-forge", shell = True )
        self.debug.installer_sync( "conda install -y -q -c conda-forge conda-pack", shell = True )


if __name__ == "__main__":
    condainit = CondaInitializer()

"""## 2. モジュール定義"""

# @title a. QuickViper2 定義
import os
import re
import json
# import yaml
import shutil
import subprocess

from pathlib import Path


get_ipython().system( "pip install an_DebugHelper" )
from an_debughelper import DebugHelper
get_ipython().system( "pip install an_UniqueEnvVar" )
from an_uniqueenvvar import UniqueEnvVar


class QuickViper2:
    def __init__( self, venv_name, python_version = "3.10" ):
        from google.colab import drive
        drive.mount( '/content/drive' )

        self.debug = DebugHelper( instance_name = "QuickViper2" )
        self.debug.enable_log_to_file_stdout()
        self.debug.enable_log_to_file_stderr()
        self.debug.enable_timestamp()

        self.python_version = python_version

    def create(self, target):
        """
        指定されたパスに conda 仮想環境を作成する。
        args:
            target( str ): 仮想環境を作成するパス
        """
        arg_target = str( target )
        self.debug.log_step(f"Creating conda environment at: { arg_target }", success = None )
        cmd = [
            "conda", "create", "-y", "-q", "-p", arg_target,
            f"python={self.python_version}"
        ]
        # result = subprocess.run(cmd, capture_output=True, text=True)
        result = self.debug.installer_sync( cmd, shell = False )
        if result.returncode != 0:
            raise Exception("Error creating conda environment:")
        else:
            self.debug.log_step("Conda environment created successfully.", success = True )


    def activate( self, target ):
        """
        展開した仮想環境をアクティブにします。
        args:
            target( str ) : アクティブにする仮想環境のパス
        """
        arg_target = str( target )
        self.debug.log_step( "activate", success = None )
        os.chdir( arg_target )
        result = self.debug.installer( command = f'bash -c "source /usr/local/etc/profile.d/conda.sh && conda activate { arg_target } && exec bash"', shell = True, back = 0 )
        if result.returncode == 0:
           self.debug.log_step(f"環境がアクティブになりました。", char="🌟")


    def pack( self, archive, target, destination ):
        """
        conda-pack を使用して、仮想環境を tar.gz 形式のアーカイブに圧縮します。
        圧縮対象の環境パスは target, 展開先はdestination,
        生成した圧縮仮想環境ファイルは archive に保存されます。
        return:
            subprocess.CompletedProcess
        """
        self.debug.log_step("圧縮処理を開始します。", success = None)

        arg_archive = str( archive )
        arg_target = str( target )
        arg_destination = str( destination )
        self.debug.log_step("圧縮処理を開始します。", success = None)

        # 既存のアーカイブがあれば削除する
        if os.path.exists( arg_archive ):
            self.debug.log_step(f"{ arg_archive } に既存のアーカイブが存在します。削除します。", success=None)
            try:
                os.remove( arg_archive )
                self.debug.log_step("既存のアーカイブを正常に削除しました。", success = True)
            except Exception as e:
                self.debug.log_step(f"既存アーカイブの削除に失敗しました: {e}", success = False)
                raise

        # conda-pack コマンドの作成
        cmd = [
            "conda-pack",
            "--prefix", arg_target,
            "--dest-prefix", arg_destination,
            "--format", "tar.gz",
            "--output", arg_archive,
        ]
        self.debug.log_step(f"conda-pack コマンド: {' '.join(cmd)}", success=None)

        # コマンド実行
        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            # result = self.debug.installer_sync( command = cmd, shell = True )
            self.debug.log_step(f"圧縮成功。アーカイブは { arg_archive} に作成されました。", success=True)
        except subprocess.CalledProcessError as e:
            self.debug.log_step(f"conda-pack が失敗しました。エラーコード: {e.returncode}", success=False)
            self.debug.log_step(f"エラーメッセージ: {e.stderr}", success=False)
            raise
        except Exception as e:
            self.debug.log_step(f"予期しないエラーが発生しました: {e}", success=False)
            raise

        self.debug.log_step("圧縮処理が完了しました。", success=True)
        return result


    def extract( self, archive, destination ):
        """
        圧縮仮想環境ファイルを展開します。
        args:
            archive( str ): 圧縮仮想環境ファイルのパス
            destination( str ): 展開先ディレクトリ
        """
        arg_archive = str( archive )
        arg_destination = str( destination )
        self.debug.log_step( "extract", success = None )
        os.makedirs( arg_destination, exist_ok = True )
        self.debug.installer( command = f"tar -xf { arg_archive } -C { arg_destination }", shell = True, back = 1 )

    # def list_archive( self ):
    #     """
    #     仮想環境のリストを作成します。
    #     """
    #     self.debug.log_step( "list_archive", success = None )
    #     get_ipython().system( f"tar tvf { self.arc_path }" )
    #     self.debug.installer( command = [ "tar", "tvf", str(self.arc_path ) ] )

    def delete_venv( self, target ):
        """
        仮想環境を消去します
        args:
            target( str ): 消去する仮想環境のパス
        """
        self.debug.log_step( "delete_venv", success = None )
        arg_target = str( target )
        if os.path.exists( arg_target ):
            shutil.rmtree( arg_target )


    def set_executable( self, target ):
        """
        実行権限を修正します。
        args:
            target( str ): 実行権限を修正する仮想環境のパス
        """
        self.debug.log_step( "set_executable", success = None )
        arg_target = Path( target ).joinpath( "bin" )
        python3_path =  arg_target / "python3.10"
        python_link = arg_target / "python"
        os.chmod(python3_path, 0o755)
        os.chmod(python_link, 0o755)
        self.debug.installer( command = f"ls -l { python3_path }", shell = True)
        self.debug.installer( command = f"ls -l { python_link }", shell = True)


    def set_env_var( self, target ):
        """
        作成した環境内の環境変数を定義する。
        args:
            target( str ): 利用する仮想環境のパス
        """
        arg_target = str ( Path( target ).joinpath( "bin" ) )
        arg_library = str( Path( target ).joinpath( "lib" ) )
        ue_path = UniqueEnvVar( "PATH" )
        ue_ld_library_path = UniqueEnvVar( "LD_LIBRARY_PATH" )
        ue_path.add_value( arg_target )
        ue_ld_library_path.add_value( arg_library )
        self.debug.log_step( f"PATH : { ue_path.get_values() }" )
        self.debug.log_step( f"LD_LIBRARY_PATH : { ue_ld_library_path.get_values() }" )


    def unlock( self, archive, destination, force_unlock = False ):
        """
        圧縮仮想環境ファイルを展開します。
        args:
            force_unlock( boolean ): Trueのとき、強制的に展開します。
        """
        self.debug.log_step( "unlock", success = None )
        arg_archive = str( archive )
        arg_destination = str( destination )
        if force_unlock and Path( arg_destination ).exists():
            shutil.rmtree( arg_destination )
        if not Path( arg_destination ).exists():
            self.extract( archive = arg_archive, destination = arg_destination )
            # self.conda_unpack()
            self.activate( target = arg_destination )
            self.set_executable( target = arg_destination )
            self.set_env_var( target = arg_destination )

    def executer( self, target, cmdlist = [ "pip", "--version" ] ):
        """
        仮想環境でプログラムを実行します。
        args:
            target( str ): コマンドを実行する仮想環境のフォルダ
            cmdlist( str [] ): 実行するプログラムのコマンドライン
        """
        self.debug.log_step( "exe_python", success = None )
        arg_target = str( Path( target ).joinpath( "bin" ) )
        os.chdir( arg_target )
        # 現在の環境変数をコピー
        env = os.environ.copy()
        # シェルを使って "env | grep LD_LIBRARY_PATH" を実行し、出力をキャプチャ
        # result = subprocess.run(["env"], env=env, capture_output=True, text=True)
        result = subprocess.run("env | grep ^LD_LIBRARY_PATH=", shell=True, env=env, capture_output=True, text=True)
        self.debug.log_step( f"result : { result }", success = None )
        # 実行コマンドをリスト形式で定義
        cmd = [ str( Path( arg_target ).joinpath( "python" ) ), "-m" ]
        full_cmd = cmd + cmdlist
        self.debug.installer( command = full_cmd )
        ### self.debug.run.commandが使えない理由はわからない。これも仮想環境絡みか?
        ### result, stdout, stderr = self.debug.run_command( command = full_cmd, stdout = True, stderr = True )
        ### self.debug.log_step( f"stdout:\n{ stdout }", success = None)
        ### self.debug.log_step( f"stderr:\n{ stdout }", success = None)

get_ipython().system( "pip install an_EasyVen" )
from an_easyven import EasyVen


# @title b. Mainloop 定義
class MainLoop():
    def __init__( self, venv_name = "kohya_env" ):
        """
        メインループ
        """
        self.debug = DebugHelper( instance_name = "MainLoop" )
        self.venv_name = venv_name
        self.even = EasyVen()
        [ self.bch_path,
          self.dst_path,
          self.cur_path,
          self.src_path,
          self.arc_path,
          self.lcl_path,
          self.bin_path,
          self.lib_path,
          self.tmp_path,
          self.org_path ] = self.even.setup( venv_name = self.venv_name )

        self.quickviper = QuickViper2( venv_name = self.venv_name )

    def remove_temp_env(self):
        """
        既存の環境があれば削除
        """
        self.quickviper.delete_venv( target = self.tmp_path )


    def create_temp_env(self):
        """
        一時的な環境を作成
        """
        # クラスを使って環境の作成、conda-pack のインストール、conda-unpack の確認を実施
        self.quickviper.create( target = self.tmp_path )

    def pack_temp_env(self):
        """
        一時的な環境を圧縮
        """
        self.quickviper.pack( archive = self.org_path, target = self.tmp_path, destination = self.lcl_path )

    def unpack_temp_env(self):
        """
        一時的な環境を展開
        """
        self.quickviper.unlock( archive = self.org_path, destination = self.lcl_path, force_unlock = True )


    def executer( self ):
        """
        pipをテスト実行
        """
        self.quickviper.executer( target = self.lcl_path )

"""# 3. テスト実行"""

if __name__ == "__main__":
    mainloop = MainLoop()
    mainloop.remove_temp_env()
    mainloop.create_temp_env()
    mainloop.pack_temp_env()
    mainloop.unpack_temp_env()
    mainloop.executer()