import os
import sys
import subprocess
import importlib.util

class LWMSetup:
    def __init__(self, repo_url="https://huggingface.co/sadjadalikhani/LWM", clone_dir="./LWM"):
        self.repo_url = repo_url
        self.clone_dir = clone_dir

    def setup(self):
        # Step 1: Clone the public repository
        self.clone_repo()

        # Step 2: Add the cloned directory to the Python path so all imports work
        sys.path.append(self.clone_dir)

        # Step 3: Dynamically import all required modules
        self.import_modules()

    def clone_repo(self):
        if not os.path.exists(self.clone_dir):
            print(f"Cloning repository from {self.repo_url} into {self.clone_dir}")

            # Clone the repository without needing an authentication token
            result = subprocess.run(["git", "clone", self.repo_url, self.clone_dir], capture_output=True, text=True)

            if result.returncode != 0:
                print(f"Error cloning the repository: {result.stderr}")
                sys.exit(1)  # Exit if cloning fails
            else:
                print(f"Repository cloned successfully into {self.clone_dir}")
        else:
            print(f"Repository already cloned into {self.clone_dir}")

    def import_modules(self):
        def import_module_from_file(module_name, file_path):
            spec = importlib.util.spec_from_file_location(module_name, file_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            return module

        try:
            self.lwm_model = import_module_from_file("lwm_model", os.path.join(self.clone_dir, "lwm_model.py"))
            self.inference = import_module_from_file("inference", os.path.join(self.clone_dir, "inference.py"))
            self.load_data = import_module_from_file("load_data", os.path.join(self.clone_dir, "load_data.py"))
            self.input_preprocess = import_module_from_file("input_preprocess", os.path.join(self.clone_dir, "input_preprocess.py"))
            print("Modules imported successfully.")
        except ModuleNotFoundError as e:
            print(f"Error importing modules: {e}")
            sys.exit(1)