"""
Tests for model conversion functionality
"""

import unittest
import sys
import os
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

TEST_MODEL_DIR = Path(__file__).parent / "model"
TEST_OUTPUT_DIR = Path(__file__).parent / "output" / "convert"
TEST_CONVERT_ONNX = TEST_MODEL_DIR / "convert" / "test_convert.onnx"


def tools_available():
    """Check if MindSpore Lite tools are installed"""
    from ohos_model_claw.tool_manager import tool_manager
    return tool_manager.check_tool_exists() and tool_manager.check_lib_dirs_exist()


class TestModelConvert(unittest.TestCase):
    """Test model conversion functionality"""

    @classmethod
    def setUpClass(cls):
        cls.output_dir = TEST_OUTPUT_DIR
        os.makedirs(cls.output_dir, exist_ok=True)

    @classmethod
    def tearDownClass(cls):
        pass

    def _get_output_file_path(self, onnx_path, output_dir):
        """Local helper to compute output .ms file path"""
        onnx_name = Path(onnx_path).stem
        return str(Path(output_dir) / f"{onnx_name}.ms")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_onnx_to_ms(self):
        """Test basic ONNX to MS conversion"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_basic")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    def test_convert_nonexistent_file(self):
        """Test error handling for nonexistent ONNX file"""
        import subprocess
        
        nonexistent_path = str(Path(__file__).parent / "nonexistent_model.onnx")
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", nonexistent_path]
        
        result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace')
        
        has_error = ("ERROR" in (result.stdout or "")) or ("Error" in (result.stderr or "")) or ("Conversion failed" in (result.stdout or "")) or (result.returncode != 0)
        self.assertTrue(has_error, "Expected error for nonexistent ONNX file")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_output_directory_creation(self):
        """Test that output directory is created if it doesn't exist"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_auto_created_dir" / "nested")
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue(Path(output_dir).exists(), "Output directory was not created")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_fp16(self):
        """Test ONNX to MS conversion with FP16 enabled"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_fp16")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir, "--fp16", "on"]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_inputDataFormat(self):
        """Test ONNX to MS conversion with inputDataFormat=NCHW"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_nchw")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir, "--inputDataFormat", "NCHW"]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_outputDataFormat(self):
        """Test ONNX to MS conversion with outputDataFormat=NCHW"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_output_nchw")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir, "--outputDataFormat", "NCHW"]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_inputShape(self):
        """Test ONNX to MS conversion with inputShape parameter"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_inputshape")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir, "--inputShape", "input:1,3,224,224"]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_optimize(self):
        """Test ONNX to MS conversion with optimize parameter"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_optimize")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [sys.executable, str(script_path), "convert", onnx_path, "--output_dir", output_dir, "--optimize", "general"]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")

    @unittest.skipIf(not tools_available(), "MindSpore Lite tools not installed")
    def test_convert_with_all_parameters(self):
        """Test ONNX to MS conversion with all new parameters combined"""
        self.assertTrue(TEST_CONVERT_ONNX.exists(), f"Test model not found: {TEST_CONVERT_ONNX}")
        
        import subprocess
        from ohos_model_claw.tool_manager import tool_manager
        
        tool_manager.apply_env_vars()
        
        output_dir = str(self.output_dir / "test_convert_all_params")
        os.makedirs(output_dir, exist_ok=True)
        
        onnx_path = str(TEST_CONVERT_ONNX)
        ms_path = self._get_output_file_path(onnx_path, output_dir)
        
        env = os.environ.copy()
        env["CONVERTER_LITE_PATH"] = str(tool_manager.converter_path)
        env["PACKAGE_ROOT_PATH"] = str(tool_manager.package_root)
        
        script_path = Path(__file__).parent.parent / "ohos_model_claw" / "ohos_model_claw.py"
        cmd = [
            sys.executable, str(script_path), "convert", onnx_path,
            "--output_dir", output_dir,
            "--fp16", "on",
            "--inputDataFormat", "NCHW",
            "--outputDataFormat", "NCHW",
            "--inputShape", "input:1,3,224,224",
            "--optimize", "general"
        ]
        
        result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120, encoding='utf-8', errors='replace')
        
        self.assertTrue("Success" in (result.stdout or "") or "OK" in (result.stdout or ""),
                        f"Expected 'Success' or 'OK' in output. stdout: {result.stdout}, stderr: {result.stderr}")
        self.assertTrue(Path(ms_path).exists(), f"Output .ms file not created: {ms_path}")


if __name__ == "__main__":
    unittest.main()