utils/disk_monitor.py

"""
ディスク使用量監視ユーティリティ
"""
import os
import time
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional
from utils.cache_manager import CacheManager


class DiskMonitor:
    """ディスク使用量を監視するクラス"""
    
    def __init__(self, warning_threshold: float = 85.0, critical_threshold: float = 90.0):
        """
        ディスク監視を初期化
        
        Args:
            warning_threshold: 警告閾値（％）
            critical_threshold: 危険閾値（％）
        """
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self.cache_manager = CacheManager()
        self.logger = logging.getLogger(__name__)
        
        # 最後の警告時刻を記録（重複警告を防ぐため）
        self.last_warning_time = None
        self.last_critical_time = None
        
    def get_status(self) -> Dict[str, Any]:
        """
        現在のディスク使用状況を取得
        
        Returns:
            ディスク使用状況の辞書
        """
        try:
            disk_usage = self.cache_manager.get_disk_usage()
            
            if "error" in disk_usage:
                return {"error": disk_usage["error"]}
            
            usage_percent = disk_usage["working_directory"]["usage_percent"]
            
            # 警告レベルを判定
            if usage_percent >= self.critical_threshold:
                level = "CRITICAL"
                action_needed = "immediate_cleanup"
            elif usage_percent >= self.warning_threshold:
                level = "WARNING"
                action_needed = "cleanup_recommended"
            else:
                level = "OK"
                action_needed = "none"
            
            return {
                "level": level,
                "usage_percent": usage_percent,
                "action_needed": action_needed,
                "disk_usage": disk_usage,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"ディスク状況の取得に失敗: {e}")
            return {"error": str(e)}
    
    def check_and_alert(self) -> Optional[Dict[str, Any]]:
        """
        ディスク使用量をチェックし、必要に応じて警告を発行
        
        Returns:
            警告が発行された場合の情報
        """
        try:
            status = self.get_status()
            
            if "error" in status:
                return None
            
            current_time = time.time()
            usage_percent = status["usage_percent"]
            
            # 危険レベルの処理
            if usage_percent >= self.critical_threshold:
                # 前回の警告から1時間以上経過している場合のみ警告
                if (self.last_critical_time is None or 
                    current_time - self.last_critical_time > 3600):
                    
                    self.last_critical_time = current_time
                    
                    # 緊急クリーンアップを実行
                    cleanup_result = self.cache_manager.emergency_cleanup()
                    
                    alert_message = {
                        "level": "CRITICAL",
                        "message": f"ディスク使用率が危険レベルに達しました: {usage_percent:.1f}%",
                        "action_taken": "emergency_cleanup",
                        "cleanup_result": cleanup_result,
                        "timestamp": datetime.now().isoformat()
                    }
                    
                    self.logger.critical(alert_message["message"])
                    print(f"[CRITICAL] {alert_message['message']}")
                    
                    return alert_message
                    
            # 警告レベルの処理
            elif usage_percent >= self.warning_threshold:
                # 前回の警告から30分以上経過している場合のみ警告
                if (self.last_warning_time is None or 
                    current_time - self.last_warning_time > 1800):
                    
                    self.last_warning_time = current_time
                    
                    # 通常のクリーンアップを実行
                    cleanup_result = self.cache_manager.cleanup_cache()
                    
                    alert_message = {
                        "level": "WARNING",
                        "message": f"ディスク使用率が警告レベルに達しました: {usage_percent:.1f}%",
                        "action_taken": "cleanup",
                        "cleanup_result": cleanup_result,
                        "timestamp": datetime.now().isoformat()
                    }
                    
                    self.logger.warning(alert_message["message"])
                    print(f"[WARNING] {alert_message['message']}")
                    
                    return alert_message
                    
            return None
            
        except Exception as e:
            self.logger.error(f"ディスク監視でエラーが発生: {e}")
            return None
    
    def continuous_monitor(self, interval: int = 300):
        """
        継続的にディスク使用量を監視
        
        Args:
            interval: 監視間隔（秒）
        """
        print(f"[INFO] ディスク監視を開始しました（間隔: {interval}秒）")
        
        try:
            while True:
                alert = self.check_and_alert()
                
                if alert:
                    print(f"[ALERT] {alert['message']}")
                
                time.sleep(interval)
                
        except KeyboardInterrupt:
            print("[INFO] ディスク監視を停止しました")
        except Exception as e:
            self.logger.error(f"ディスク監視中にエラーが発生: {e}")
            print(f"[ERROR] ディスク監視中にエラーが発生: {e}")
    
    def generate_report(self) -> str:
        """
        現在の状況レポートを生成
        
        Returns:
            レポート文字列
        """
        try:
            status = self.get_status()
            
            if "error" in status:
                return f"レポート生成エラー: {status['error']}"
            
            disk_usage = status["disk_usage"]
            working_dir = disk_usage["working_directory"]
            cache_info = disk_usage["cache"]
            
            report = f"""
=== ディスク使用量レポート ===
生成日時: {status['timestamp']}
警告レベル: {status['level']}

【作業用ディレクトリ】
パス: {working_dir['path']}
使用率: {working_dir['usage_percent']:.1f}%
使用量: {working_dir['used_gb']:.2f}GB / {working_dir['total_gb']:.2f}GB
空き容量: {working_dir['free_gb']:.2f}GB

【Huggingfaceキャッシュ】
キャッシュサイズ: {cache_info.get('total_size_gb', 0):.2f}GB
リポジトリ数: {cache_info.get('num_repos', 0)}
キャッシュパス: {cache_info.get('cache_path', 'N/A')}

【推奨アクション】
"""
            
            if status["action_needed"] == "immediate_cleanup":
                report += "⚠️  緊急: 即座にキャッシュクリーンアップが必要です"
            elif status["action_needed"] == "cleanup_recommended":
                report += "⚠️  推奨: キャッシュクリーンアップを検討してください"
            else:
                report += "✅ 現在のところ、アクションは不要です"
            
            return report
            
        except Exception as e:
            return f"レポート生成中にエラーが発生: {e}"


def main():
    """メイン関数（スタンドアロン実行用）"""
    import argparse
    
    parser = argparse.ArgumentParser(description="ディスク使用量監視ツール")
    parser.add_argument("--monitor", action="store_true", help="継続監視を開始")
    parser.add_argument("--report", action="store_true", help="現在の状況レポートを表示")
    parser.add_argument("--interval", type=int, default=300, help="監視間隔（秒）")
    parser.add_argument("--warning", type=float, default=85.0, help="警告閾値（％）")
    parser.add_argument("--critical", type=float, default=90.0, help="危険閾値（％）")
    
    args = parser.parse_args()
    
    monitor = DiskMonitor(
        warning_threshold=args.warning,
        critical_threshold=args.critical
    )
    
    if args.monitor:
        monitor.continuous_monitor(interval=args.interval)
    elif args.report:
        print(monitor.generate_report())
    else:
        # デフォルト: 1回だけチェック
        status = monitor.get_status()
        if "error" in status:
            print(f"エラー: {status['error']}")
        else:
            print(f"ディスク使用率: {status['usage_percent']:.1f}% ({status['level']})")
            alert = monitor.check_and_alert()
            if alert:
                print(f"警告: {alert['message']}")


if __name__ == "__main__":
    main()