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

Automatically generated by Colab.

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

# QuickViper3 更新履歴
QuickViper2を元に、ボトルネックになっていた仮想環境圧縮・展開プロセスをcondaに一任した。


```
2025/04/25 1.0.0 完成
2025/04/25 1.1.2 docscriptを記入したり、いろいろした。
2025/04/26 1.1.3 extractの引数修正など
2025/04/27 1.2.0 全体動作の安定化
```

# QuickViper3

## 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. QuickViper3 定義
import os
import json
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_toolkit" )
from an_toolkit import ToolKit
get_ipython().system( "pip install an_uniqueenvvar" )
from an_uniqueenvvar import UniqueEnvVar
get_ipython().system( "pip install an_EasyVen" )
from an_easyven import EasyVen

class QuickViper3:
    """
    A conda-based virtual environment manager that uses venv_name consistently.
    """
    def __init__(self, venv_name: str,  python_version: str = "3.10"):
        """
        QuickViper3
        Conda利用仮想環境管理システム
        args:
            venv_name( str ): 作成する仮想環境の名前
            python_version( str ): 仮想環境にインストールするPythonのバージョン
        """
        self.venv_name = venv_name
        self.python_version = python_version
        self.debug = DebugHelper(instance_name="QuickViper3")
        self.debug.enable_log_to_file_stdout()
        self.debug.enable_log_to_file_stderr()
        self.debug.enable_timestamp()
        self.toolkit = ToolKit()
        self.even = EasyVen()
        [ self.bch_path,
          self.dst_path,
          self.cur_path,
          self.src_path,
          self.arc_path,
          self.arz_path,
          self.lcl_path,
          self.bin_path,
          self.lib_path,
          self.tmp_path,
          self.org_path,
          self.orz_path,
          self.enva_path,
          self.envb_path,
          self.uni_path ] = self.even.setup( venv_name = self.venv_name )
        ## self.env_prefix = os.path.join( self.dst_path, self.venv_name, self.bch_path )

    def create(self):
        """
        インスタンス作成時に渡した名前で仮想環境を作成します
        """
        self.debug.log_step(f"Creating conda env '{self.venv_name}' with Python {self.python_version}", success=None)
        cmd = [
            "conda", "create", "-y", "-q", "-p", str( self.lcl_path ),
            f"python={self.python_version}"
        ]
        result = self.toolkit.executor(cmd=cmd, shell=False)
        if result.returncode != 0:
            raise Exception(f"Failed to create env '{self.venv_name}'.")
        self.debug.log_step(f"Environment '{self.venv_name}' created.", success=True)

    def remove( self ):
        """
        作成した仮想環境を消去します
        """
        if not self.lcl_path.exists():
            self.debug.log_step( "消去対象の仮想環境は存在しないようです。スキップします", success = True )
            return

        cmd = [
            "conda", "env", "remove", "-y", "-q", "-p", str( self.lcl_path )
        ]
        result = self.toolkit.executor(cmd=cmd, shell=False)
        if result.returncode != 0:
            raise Exception(f"Failed to remove env '{self.venv_name}'.")
        self.debug.log_step(f"Environment '{self.venv_name}' removed.", success=True)

    def create_from_environment_yml(self, yml_path: Path):
        """
        yml_pathに渡したenvironment.ymlを元に、仮想環境を構築します
        args:
            yml_path( Path ): 環境構築に使用するenvironment.ymlのパス
        """
        self.debug.log_step(f"Creating env '{ self.venv_name }' from { yml_path }", success=None)
        cmd = [
            "conda", "env", "create", "-p", str( self.lcl_path ),
            "-f", str(yml_path)
        ]
        result = self.toolkit.executor( cmd = cmd, shell = False )
        if result.returncode != 0:
            raise Exception(f"Failed to create env from { yml_path }.")
        self.debug.log_step(f"Environment '{ self.venv_name }' created from YAML.", success=True)

    def export_environment_yml(self, output_path: Path, no_builds: bool = True):
        """
        conda仮想環境の情報を、output_pathのenvironment.ymlにエクスポートします
        args:
            output_path( Path ): 環境構築に使用するenvironment.ymlのパス
            no_builds( bool ): 依存関係の更新を無効にするかどうか
        """
        self.debug.log_step(f"Exporting env '{self.venv_name}' to {output_path}", success=None)
        cmd = ["conda", "env", "export", "-p", str( self.lcl_path ) ]
        if no_builds:
            cmd.append("--no-builds")
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        yaml_text = result.stdout
        with open(output_path, "w") as f:
            f.write(yaml_text)
        self.debug.log_step(f"Exported environment.yml to {output_path}", success=True)

    def pack(self, archive: Path):
        """
        conda-packを使い、仮想環境をtar.gz形式で圧縮します。
        args:
            archive( Path ): 圧縮された仮想環境の保存先パス
        """
        self.debug.log_step("Packing environment...", success=None)
        self.toolkit.file_remover( str( archive ) )
        cmd = [
            "conda-pack",
            "--prefix", str( self.lcl_path ),
            "--output", str( archive ),
            "--format", "tar.gz",
            "--ignore-missing-files"
        ]
        self.toolkit.executor(cmd=cmd)
        self.debug.log_step(f"Environment packed to { archive }.", success=True)


    def extract(self, archive: Path ):
        """
        conda圧縮仮想環境ファイルを、自らの環境として展開します
        args:
            archive( Path ): 圧縮仮想環境フォルダへのパス
        """
        self.debug.log_step(f"Extracting { archive } to { self.lcl_path }", success=None)
        if self.lcl_path.exists():
            shutil.rmtree( self.lcl_path )
        os.makedirs( self.lcl_path, exist_ok = True )
        cmd = ["tar", "-xf", str( archive ), "-C", str( self.lcl_path )]
        self.toolkit.executor(cmd=cmd, shell=False)
        self.debug.log_step(f"Archive extracted to { self.lcl_path }", success=True)

    def activate(self):
        """
        conda仮想環境をアクティブにします
        """
        self.debug.log_step(f"Activating env '{self.venv_name}'", success=None)
        cmd = f'bash -c "source /usr/local/etc/profile.d/conda.sh && conda activate { str( self.lcl_path ) }"'
        result = self.toolkit.executor_sync( cmd = cmd, shell = True )
        if result.returncode == 0:
            self.debug.log_step(f"Environment '{self.venv_name}' activated.", success=True)

    def deactivate(self):
        """
        conda仮想環境をディアクティブにします
        """
        self.debug.log_step(f"Deactivating env", success=None)
        cmd = f'bash -c "source /usr/local/etc/profile.d/conda.sh && conda deactivate"'
        result = self.toolkit.executor_sync( cmd = cmd, shell = True )
        if result.returncode == 0:
            self.debug.log_step(f"Environment deactivated.", success=True)

    def delete(self):
        """
        展開済みの環境ディレクトリを丸ごと削除します
        存在しない場合はスキップ
        """
        self.debug.log_step("delete", success=None)
        if Path( self.lcl_path ).exists():
            self.toolkit.dir_remover( str( self.lcl_path ) )
            self.debug.log_step(f"Removed environment directory: { self.venv_name }", success=True)
        else:
            # 環境がないならエラーにせず「何もせず通過」
            self.debug.log_step(f"No environment found at: { self.venv_name }, skipping deletion.", success=True)


    def executor_sync(self, cmdlist: list[str] ):
        """
        この conda 環境内で、任意のコマンドを実行します。
        args:
            cmdlist( list[ str ] ): 実行したいコマンドと引数のリスト, 例 ["python", "-m", "pip", "install", "requests"]
        """
        # conda run を使って、名前指定で環境を呼び出し
        self.debug.log_step( "executor", success = None )
        cmd = [ "conda", "run", "-p", str( self.lcl_path ) ] + cmdlist
        result = self.toolkit.executor_sync( cmd = cmd, shell = False )
        if result.returncode != 0:
            self.debug.log_step(f"Command failed: {' '.join(cmd)}", success = False )
        return result

# @title b. Mainloop 定義

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

class MainLoop():
    def __init__( self, venv_name = "kohya_env" ):
        """
        メインループ
        """
        self.debug = DebugHelper( instance_name = "MainLoop" )
        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.arz_path,
          self.lcl_path,
          self.bin_path,
          self.lib_path,
          self.tmp_path,
          self.org_path,
          self.orz_path,
          self.enva_path,
          self.envb_path,
          self.uni_path ] = self.even.setup( venv_name = self.venv_name )
        self.quickviper = QuickViper3( venv_name = self.venv_name )

    def remove_temp_env(self):
        """
        既存の環境があれば削除
        """
        self.quickviper.deactivate()
        self.quickviper.remove()


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

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

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

    def export_to_yml( self ):
        self.quickviper.export_environment_yml( output_path = os.path.join( self.arc_path, "environment.yml" ), no_builds = True )

    def create_from_yml( self ):
        self.quickviper.create_from_environment_yml( yml_path = os.path.join( self.arc_path, "environment.yml" ) )
        self.quickviper.activate()

    def executor_sync( self ):
        """
        pipをテスト実行
        """
        self.quickviper.executor_sync( cmdlist = ["python", "-m", "pip", "list"] )

"""## 3. テスト"""

# @title a. MainLoop 実行
if __name__ == "__main__":
    mainloop = MainLoop()
    mainloop.remove_temp_env()
    mainloop.create_from_yml()
    # mainloop.remove_temp_env()
    # mainloop.pack_temp_env()
    # mainloop.unpack_temp_env()
    # mainloop.executor_sync()
    # mainloop.export_to_yml()
    # mainloop.remove_temp_env()
    # mainloop.create_from_yml()