#!/usr/bin/env python3
"""
腎臓移植データ処理システム - 分割実行バッチスクリプト（Python版）
50,000件のデータを分割して処理し、メモリ不足を回避します
"""

import argparse
import subprocess
import sys
import os
import time
from datetime import datetime
from typing import List, Tuple
import logging

def setup_logging(log_dir: str, timestamp: str) -> logging.Logger:
    """ログ設定の初期化"""
    os.makedirs(log_dir, exist_ok=True)
    
    # ログファイル設定
    log_file = os.path.join(log_dir, f'batch_execution_{timestamp}.log')
    
    # ロガー設定
    logger = logging.getLogger('batch_executor')
    logger.setLevel(logging.INFO)
    
    # ファイルハンドラ
    file_handler = logging.FileHandler(log_file, encoding='utf-8')
    file_handler.setLevel(logging.INFO)
    
    # コンソールハンドラ
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    # フォーマッター
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    return logger

def execute_batch(start: int, limit: int, optimize: str = None, logger: logging.Logger = None) -> bool:
    """単一バッチの実行"""
    cmd = [sys.executable, 'backend/jin.py', '--start', str(start), '--limit', str(limit)]
    
    if optimize:
        cmd.extend(['--optimize', optimize])
    
    logger.info(f"実行コマンド: {' '.join(cmd)}")
    
    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        if logger:
            logger.info("バッチ実行成功")
            if result.stdout:
                logger.debug(f"標準出力: {result.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        if logger:
            logger.error(f"バッチ実行失敗: {e}")
            if e.stdout:
                logger.error(f"標準出力: {e.stdout}")
            if e.stderr:
                logger.error(f"標準エラー: {e.stderr}")
        return False

def main():
    """メイン処理"""
    parser = argparse.ArgumentParser(description="腎臓移植データ分割実行バッチ（Python版）")
    parser.add_argument("-t", "--total", type=int, default=50000, help="処理総件数（デフォルト: 50000）")
    parser.add_argument("-b", "--batch", type=int, default=20000, help="バッチサイズ（デフォルト: 20000）")
    parser.add_argument("-o", "--optimize", choices=['O', 'OO'], help="最適化レベル（O または OO）")
    parser.add_argument("-l", "--log-dir", default="./batch_logs", help="ログ出力ディレクトリ（デフォルト: ./batch_logs）")
    parser.add_argument("--continue-on-error", action="store_true", help="エラー時も継続実行")
    parser.add_argument("--sleep", type=int, default=5, help="バッチ間の待機時間（秒、デフォルト: 5）")
    
    args = parser.parse_args()
    
    # タイムスタンプ
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    # ログ設定
    logger = setup_logging(args.log_dir, timestamp)
    
    # 実行開始
    start_time = datetime.now()
    logger.info("=" * 50)
    logger.info("腎臓移植データ分割実行バッチ開始（Python版）")
    logger.info("=" * 50)
    logger.info(f"開始時刻: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info(f"処理総件数: {args.total:,}")
    logger.info(f"バッチサイズ: {args.batch:,}")
    logger.info(f"最適化レベル: {args.optimize or 'なし'}")
    logger.info(f"ログディレクトリ: {args.log_dir}")
    
    # バッチ数計算
    batch_count = (args.total + args.batch - 1) // args.batch
    logger.info(f"予定バッチ数: {batch_count}")
    logger.info("=" * 50)
    
    # 実行カウンタ
    success_count = 0
    failed_count = 0
    failed_batches: List[Tuple[int, int, int]] = []
    
    # 分割実行のメインループ
    for i in range(0, args.total, args.batch):
        batch_start = i
        batch_end = min(i + args.batch, args.total)
        actual_batch_size = batch_end - batch_start
        batch_num = i // args.batch + 1
        
        logger.info(f"バッチ {batch_num}/{batch_count} 開始: {batch_start:,}-{batch_end:,} ({actual_batch_size:,} 件)")
        
        batch_start_time = datetime.now()
        
        # バッチ実行
        success = execute_batch(batch_start, actual_batch_size, args.optimize, logger)
        
        batch_end_time = datetime.now()
        batch_duration = batch_end_time - batch_start_time
        
        if success:
            logger.info(f"バッチ {batch_num}/{batch_count} 成功 (処理時間: {batch_duration})")
            success_count += 1
            
            # メモリクリーンアップのための待機
            if args.sleep > 0:
                logger.info(f"メモリクリーンアップのため {args.sleep} 秒待機中...")
                time.sleep(args.sleep)
        else:
            logger.error(f"バッチ {batch_num}/{batch_count} 失敗 (処理時間: {batch_duration})")
            failed_count += 1
            failed_batches.append((batch_num, batch_start, batch_end))
            
            if not args.continue_on_error:
                response = input("バッチが失敗しました。続行しますか？ [y/N]: ")
                if response.lower() not in ['y', 'yes']:
                    logger.info("ユーザーによる中断")
                    break
            
            # 失敗後の待機時間を長くする
            if args.sleep > 0:
                extended_sleep = args.sleep * 2
                logger.info(f"失敗後の待機時間: {extended_sleep} 秒")
                time.sleep(extended_sleep)
        
        logger.info("-" * 30)
    
    # 実行終了
    end_time = datetime.now()
    total_duration = end_time - start_time
    
    # 結果サマリー
    logger.info("=" * 50)
    logger.info("分割実行完了サマリー")
    logger.info("=" * 50)
    logger.info(f"開始時刻: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info(f"終了時刻: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info(f"総処理時間: {total_duration}")
    logger.info(f"処理総件数: {args.total:,}")
    logger.info(f"バッチサイズ: {args.batch:,}")
    logger.info(f"成功バッチ数: {success_count}")
    logger.info(f"失敗バッチ数: {failed_count}")
    
    if failed_batches:
        logger.info("")
        logger.info("失敗したバッチ:")
        for batch_num, start, end in failed_batches:
            logger.info(f"  - バッチ {batch_num}: {start:,}-{end:,}")
        
        logger.info("")
        logger.info("失敗したバッチの再実行コマンド:")
        for batch_num, start, end in failed_batches:
            limit = end - start
            retry_cmd = f"python3 backend/jin.py --start {start} --limit {limit}"
            if args.optimize:
                retry_cmd += f" --optimize {args.optimize}"
            logger.info(f"  {retry_cmd}")
    
    logger.info("=" * 50)
    
    # 終了コード
    return failed_count == 0

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)