import os
import shutil
import subprocess
import logging
import time
from build_manager import BuildManager, BuildState

class Builder:
    def __init__(self, build_id: str = None, workdir: str = None, **kwargs):
        self.logger = logging.getLogger(__name__)
        self.bm = BuildManager.get_singleton()
        
        # --- MODIFICATION START ---
        # Reason: Upstream (web app) and worker container need a consistent 
        # base directory. We default to /base where the ArduPilot source lives.
        self.workdir = workdir or os.environ.get("CBS_BASEDIR", "/base")
        # --- MODIFICATION END ---
        
        self.build_id = build_id
        if self.build_id:
            self._setup_paths()

    def _setup_paths(self):
        self.artifacts_dir = self.bm.get_build_artifacts_dir_path(self.build_id)
        self.log_file = self.bm.get_build_log_path(self.build_id)
        self.info = self.bm.get_build_info(self.build_id)

    def __run_cmd(self, cmd, cwd, log_handle):
        # --- MODIFICATION START ---
        # Reason: Using python3 explicitly to run 'waf' is more robust than 
        # relying on the execution bit (+x) inside a Docker volume.
        process = subprocess.Popen(
            cmd, cwd=cwd, stdout=log_handle, stderr=subprocess.STDOUT, text=True
        )
        process.wait()
        if process.returncode != 0:
            raise subprocess.CalledProcessError(process.returncode, cmd)
        # --- MODIFICATION END ---

    def build(self, build_id: str = None):
        if build_id:
            self.build_id = build_id
            self._setup_paths()

        if not self.build_id:
            raise ValueError("No build_id provided to Builder.")

        try:
            self.logger.info(f"[{self.build_id}] Starting build process...")
            os.makedirs(self.artifacts_dir, exist_ok=True)
            
            # --- MODIFICATION START ---
            # Reason: We must ensure we are in the directory containing 'waf'.
            # If the path is wrong, the build fails immediately.
            repo = self.workdir
            if not os.path.exists(os.path.join(repo, "waf")):
                self.logger.error(f"[{self.build_id}] waf not found in {repo}")
                raise FileNotFoundError(f"Cannot find waf in {repo}")

            with open(self.log_file, "a") as log:
                log.write(f"Starting build: {self.info.vehicle_id} on {self.info.board}\n")
                
                # --- EQUALMASS OVERLAY INJECTION ---
                # Reason: This is where your custom files from the sidecar manager 
                # are merged into the ArduPilot source before the compiler starts.
                overlay_dir = "/app/overlay"
                if os.path.exists(overlay_dir) and os.listdir(overlay_dir):
                    self.logger.info(f"[{self.build_id}] Custom overlay found. Injecting...")
                    shutil.copytree(overlay_dir, repo, dirs_exist_ok=True)
                else:
                    self.logger.info(f"[{self.build_id}] No overlay files. Building vanilla.")
                
                # Running waf via python3 for better Docker compatibility
                self.logger.info(f"[{self.build_id}] Running waf configure...")
                self.__run_cmd(["python3", "waf", "configure", "--board", self.info.board], repo, log)
                
                self.logger.info(f"[{self.build_id}] Running waf build...")
                self.__run_cmd(["python3", "waf", self.info.vehicle_id], repo, log)
                # --- MODIFICATION END ---

            self.bm.update_build_progress_state(self.build_id, BuildState.SUCCESS)
        except Exception as e:
            self.logger.error(f"[{self.build_id}] Build failed: {e}")
            self.bm.update_build_progress_state(self.build_id, BuildState.FAILURE)

    # --- MODIFICATION START ---
    # Reason: The builder container needs this loop to stay alive and poll 
    # Redis for jobs. The web app also checks for this method during startup.
    def run(self):
        """Main worker heartbeat loop."""
        self.logger.info("Worker online. Waiting for builds from Redis...")
        
        while True:
            try:
                build_id = self.bm.get_next_build_id()
                if build_id:
                    self.logger.info(f"Job received: {build_id}")
                    self.build(build_id=build_id)
                
                time.sleep(2) # Prevent CPU spiking
            except Exception as e:
                self.logger.error(f"Error in worker loop: {e}")
                time.sleep(5)
    # --- MODIFICATION END ---