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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1y18s5hZVb3wOcdyOz3e5Z3D3HXwbg6zc

# 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追加 他
```

# 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_EasyVen" )
from an_easyven import EasyVen
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.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.even.setup( venv_name = self.venv_name, back = 0 )

        self.python_version = python_version

    def create(self):
        """
        指定されたパスに conda 仮想環境を作成する。
        """
        self.debug.log_step(f"Creating conda environment at: {self.tmp_path}", success = None )
        cmd = [
            "conda", "create", "-y", "-q", "-p", str(self.tmp_path),
            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 install_conda_pack(self):
    #     """
    #     作成した環境内に conda-pack を pip 経由でインストールする。
    #     """
    #     self.debug.log_step(f"Installing conda-pack in environment: {self.tmp_path}", success = None )
    #     # conda run を用いて、対象環境内で pip install を実行
    #     cmd = [
    #         "conda", "run", "-p", str(self.tmp_path),
    #         "pip", "install", "conda-pack"
    #     ]
    #     result = subprocess.run(cmd, capture_output=True, text=True)
    #     # result = self.debug.installer_sync( cmd, shell = False )
    #     if result.returncode != 0:
    #         raise Exception("Error installing conda-pack:")
    #     else:
    #         self.debug.log_step( "conda-pack installed successfully.", success = True )

    # def install_conda_pack(self):
    #     self.debug.log_step(f"Installing conda-pack in environment: {self.tmp_path}", success=None)
    #     cmd = [
    #         "conda", "install", "-y", "-q", "-p", str(self.tmp_path),
    #         "-c", "conda-forge", "conda-pack"
    #     ]
    #     result = subprocess.run(cmd, capture_output=True, text=True)
    #     if result.returncode != 0:
    #         raise Exception("Error installing conda-pack:")
    #     else:
    #         self.debug.log_step("conda-pack installed successfully.", success=True)

    def activate( self ):
        """
        展開した仮想環境を有効にします。
        """
        self.debug.log_step( "activate", success = None )
        os.chdir( self.lcl_path )
        # get_ipython().system( f"./bin/conda-unpack" )
        # cmd = f"{ Path( self.lcl_path ).joinpath( self.bin_path ) }/conda-unpack"
        # result = self.debug.installer( command = f'{ Path( self.lcl_path ).joinpath( self.bin_path ) }/conda-unpack', shell = True, back = 1 )
        result = self.debug.installer( command = f'bash -c "source /usr/local/etc/profile.d/conda.sh && conda activate { self.lcl_path } && exec bash"', shell = True, back = 0 )
        if result.returncode == 0:
           self.debug.log_step(f"環境が有効になりました。", char="🌟")


    def check_conda_unpack(self):
        """
        作成した環境の bin ディレクトリに conda-unpack があるか確認する。
        """
        conda_unpack_path = self.tmp_path / self.bin_path / "conda-unpack"
        if conda_unpack_path.exists():
            self.debug.log_step( f"conda-unpack is available at:{ conda_unpack_path }", success = True )
        else:
            self.debug.log_step( "conda-unpack NOT found in the environment.", success = False )

    def pack( self ):
        """
        conda-pack を使用して、仮想環境を tar.gz 形式のアーカイブに圧縮します。
        圧縮対象の環境パスは self.tmp_path, 展開先はself.lcl_path,
        生成した圧縮仮想環境ファイルは self.arc_path に保存されます。
        また、--include conda-unpack オプションを付けて、conda-unpack スクリプトも同梱します。
        return:
            subprocess.CompletedProcess
        """
        self.debug.log_step("圧縮処理を開始します。", success = None)

        # 既存のアーカイブがあれば削除する
        if os.path.exists( self.arc_path ):
            self.debug.log_step(f"{ self.arc_path } に既存のアーカイブが存在します。削除します。", success=None)
            try:
                os.remove( self.arc_path )
                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", str( self.tmp_path ),
            "--dest-prefix", str( self.lcl_path ),
            "--format", "tar.gz",
            "--include", "conda-unpack",
            "--output", str(self.arc_path),
            "--ignore-missing-files"
        ]
        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"圧縮成功。アーカイブは {self.arc_path} に作成されました。", 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 ):
        """
        圧縮仮想環境ファイルを展開します。
        """
        self.debug.log_step( "extract", success = None )
        os.chdir( self.cur_path )
        os.makedirs( self.lcl_path, exist_ok = True )
        self.debug.installer( command = f"tar -xf { self.arc_path } -C {self.lcl_path}", 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 ):
        """
        仮想環境を消去します
        """
        self.debug.log_step( "delete_venv", success = None )
        shutil.rmtree( self.lcl_path )


    def set_executable( self ):
        """
        実行権限を修正します。
        """
        self.debug.log_step( "set_executable", success = None )
        python3_path = self.bin_path / "python3.10"
        python_link = self.bin_path / "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 ):
        """
        作成した環境内の環境変数を定義する。
        """
        ue_path = UniqueEnvVar( "PATH" )
        ue_ld_library_path = UniqueEnvVar( "LD_LIBRARY_PATH" )
        ue_path.add_value( str( self.bin_path ) )
        ue_ld_library_path.add_value( str( self.lib_path ) )
        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, force_unlock = False ):
        """
        圧縮仮想環境ファイルを展開します。
        args:
            force_unlock( boolean ): Trueのとき、強制的に展開します。
        """
        self.debug.log_step( "unlock", success = None )
        if force_unlock and Path( self.lcl_path ).exists():
            shutil.rmtree( self.lcl_path )
        if not Path( self.lcl_path ).exists():
            self.extract()
            # self.conda_unpack()
            self.activate()
            self.set_executable()
            self.set_env_var()

    def executer( self, cmdlist = [ "pip", "--version" ] ):
        """
        仮想環境でプログラムを実行します。
        args:
            cmdlist( str [] ): 実行するプログラムのコマンドライン
        """
        self.debug.log_step( "exe_python", success = None )
        os.chdir( self.bin_path )
        # 現在の環境変数をコピー
        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(self.bin_path / "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)

# @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.even.setup( venv_name = self.venv_name )

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

    def remove_temp_env(self):
        """
        既存の環境があれば削除
        """
        if self.tmp_path.exists():
            print(f"Removing existing environment at: { self.tmp_path }")
            subprocess.run(["rm", "-rf", str( self.tmp_path )])

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

    def pack_temp_env(self):
        """
        一時的な環境を圧縮
        """
        self.quickviper.pack()

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


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

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

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