# scripts/base_runner.py
import subprocess
import os
from typing import Dict, Any, List, Tuple

class BaseRunner:
    """基础模型运行器"""
    
    def __init__(self, config: Dict[str, Any], env):
        self.config = config
        self.env = env
        self.base_config = self.config.get('base', {})
        self.common_config = self.base_config.get('common', {})
        
    def build_command(self, model_file: str, model_config: Dict[str, Any], 
                     batch_size: int) -> List[str]:
        """构建migraphx-driver命令"""
        
        # 基础命令
        cmd = [
            self.common_config.get('migraphx_driver', '/opt/dtk/bin/migraphx-driver'),
            "perf",
            model_file
        ]
        
        # FP16选项
        if self.common_config.get('fp16', True):
            cmd.append("--fp16")
        
        # 迭代次数
        iterations = model_config.get('iterations', self.common_config.get('iterations', 100))
        cmd.extend(["-n", str(iterations)])
        
        # 输入维度
        inputs = model_config.get('inputs', [])
        for input_config in inputs:
            input_name = input_config.get('name', 'x')
            shape = input_config.get('shape', [])
            
            # 将batch size插入到形状的第一个位置
            full_shape = [str(batch_size)] + [str(dim) for dim in shape]
            
            cmd.extend(["--input-dim", f"@{input_name}", *full_shape])
        
        # 额外参数
        extra_args = model_config.get('extra_args', [])
        if extra_args:
            cmd.extend(extra_args)
            
        return cmd
    
    def run_model(self, cmd: List[str], log_file: str) -> Tuple[bool, str]:
        """运行模型并记录日志"""
        
        try:
            # 执行命令并同时输出到终端和文件
            with open(log_file, 'w') as f:
                print(f"执行命令: {' '.join(cmd)}")
                
                # 运行命令，实时输出到终端和文件
                process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                # 实时读取输出
                output_lines = []
                for line in process.stdout:
                    print(line, end='')
                    f.write(line)
                    output_lines.append(line)
                
                process.wait()
                
                output = ''.join(output_lines)
                success = (process.returncode == 0)
                
                return success, output
                
        except Exception as e:
            error_msg = f"执行命令失败: {e}"
            print(error_msg)
            with open(log_file, 'a') as f:
                f.write(f"\n错误: {error_msg}\n")
            return False, error_msg