from __future__ import annotations

import os
import shutil
import shlex
import subprocess
import time
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from edgeeda.utils import ensure_dir


@dataclass
class RunResult:
    return_code: int
    runtime_sec: float
    cmd: str
    stdout: str
    stderr: str
    
    def is_success(self) -> bool:
        """Check if the run was successful."""
        return self.return_code == 0
    
    def error_summary(self, max_lines: int = 5) -> str:
        """Extract key error information from stderr."""
        if self.is_success():
            return "Success"
        
        lines = self.stderr.split('\n')
        # Find lines with error keywords
        error_lines = [
            l for l in lines 
            if any(kw in l.lower() for kw in ['error', 'fatal', 'failed', 'exception'])
        ]
        
        if error_lines:
            return '\n'.join(error_lines[-max_lines:])
        
        # Fallback: last few lines of stderr
        if lines:
            return '\n'.join(lines[-max_lines:])
        
        return f"Command failed with return code {self.return_code}"


class ORFSRunner:
    """
    Minimal ORFS interface:
      - Runs `make <target> DESIGN_CONFIG=... FLOW_VARIANT=... VAR=...`
      - Uses ORFS_FLOW_DIR (OpenROAD-flow-scripts/flow) as working directory.
    """

    def __init__(self, orfs_flow_dir: str):
        self.flow_dir = os.path.abspath(orfs_flow_dir)
        if not os.path.isdir(self.flow_dir):
            raise FileNotFoundError(f"ORFS flow dir not found: {self.flow_dir}")
        self._openroad_fallback = os.path.abspath(
            os.path.join(self.flow_dir, "..", "tools", "install", "OpenROAD", "bin", "openroad")
        )
        self._opensta_fallback = os.path.abspath(
            os.path.join(self.flow_dir, "..", "tools", "install", "OpenROAD", "bin", "sta")
        )
        self._yosys_fallback = os.path.abspath(
            os.path.join(self.flow_dir, "..", "tools", "install", "yosys", "bin", "yosys")
        )

    def _build_env(self) -> Dict[str, str]:
        env = os.environ.copy()
        openroad_exe = env.get("OPENROAD_EXE")
        if not openroad_exe or not os.path.isfile(openroad_exe) or not os.access(openroad_exe, os.X_OK):
            if os.path.isfile(self._openroad_fallback) and os.access(self._openroad_fallback, os.X_OK):
                env["OPENROAD_EXE"] = self._openroad_fallback
            else:
                found = shutil.which("openroad")
                if found:
                    env["OPENROAD_EXE"] = found
        opensta_exe = env.get("OPENSTA_EXE")
        if not opensta_exe or not os.path.isfile(opensta_exe) or not os.access(opensta_exe, os.X_OK):
            if os.path.isfile(self._opensta_fallback) and os.access(self._opensta_fallback, os.X_OK):
                env["OPENSTA_EXE"] = self._opensta_fallback
            else:
                found = shutil.which("sta")
                if found:
                    env["OPENSTA_EXE"] = found
        yosys_exe = env.get("YOSYS_EXE")
        if not yosys_exe or not os.path.isfile(yosys_exe) or not os.access(yosys_exe, os.X_OK):
            if os.path.isfile(self._yosys_fallback) and os.access(self._yosys_fallback, os.X_OK):
                env["YOSYS_EXE"] = self._yosys_fallback
            else:
                found = shutil.which("yosys")
                if found:
                    env["YOSYS_EXE"] = found
        return env

    def run_make(
        self,
        target: str,
        design_config: str,
        flow_variant: str,
        overrides: Dict[str, str],
        timeout_sec: Optional[int] = None,
        extra_make_args: Optional[List[str]] = None,
        max_retries: int = 0,
    ) -> RunResult:
        """
        Run make command with optional retry logic.
        
        Args:
            target: Make target (e.g., 'synth', 'place', 'route')
            design_config: Design configuration path
            flow_variant: Flow variant identifier
            overrides: Dictionary of make variable overrides
            timeout_sec: Timeout in seconds
            extra_make_args: Additional make arguments
            max_retries: Maximum number of retries for transient failures
        
        Returns:
            RunResult with command execution details
        """
        extra_make_args = extra_make_args or []
        # Build make command
        cmd_list = [
            "make",
            target,
            f"DESIGN_CONFIG={design_config}",
            f"FLOW_VARIANT={flow_variant}",
        ]
        for k, v in overrides.items():
            cmd_list.append(f"{k}={v}")
        cmd_list += extra_make_args

        cmd_str = " ".join(shlex.quote(x) for x in cmd_list)

        # Retry logic for transient failures
        last_result = None
        for attempt in range(max_retries + 1):
            t0 = time.time()
            try:
                env = self._build_env()
                p = subprocess.run(
                    cmd_list,
                    cwd=self.flow_dir,
                    capture_output=True,
                    text=True,
                    timeout=timeout_sec,
                    env=env,
                )
                dt = time.time() - t0
                result = RunResult(
                    return_code=p.returncode,
                    runtime_sec=dt,
                    cmd=cmd_str,
                    stdout=p.stdout[-20000:],  # keep tail
                    stderr=p.stderr[-20000:],
                )
                
                # If successful or no more retries, return
                if result.is_success() or attempt >= max_retries:
                    return result
                
                last_result = result
                
                # Exponential backoff before retry
                if attempt < max_retries:
                    wait_time = 2 ** attempt
                    time.sleep(wait_time)
                    
            except subprocess.TimeoutExpired:
                dt = time.time() - t0
                result = RunResult(
                    return_code=124,  # Standard timeout exit code
                    runtime_sec=dt,
                    cmd=cmd_str,
                    stdout="",
                    stderr=f"Command timed out after {timeout_sec} seconds",
                )
                if attempt >= max_retries:
                    return result
                last_result = result
                if attempt < max_retries:
                    time.sleep(2 ** attempt)
            except Exception as e:
                dt = time.time() - t0
                result = RunResult(
                    return_code=1,
                    runtime_sec=dt,
                    cmd=cmd_str,
                    stdout="",
                    stderr=f"Exception during execution: {str(e)}",
                )
                if attempt >= max_retries:
                    return result
                last_result = result
                if attempt < max_retries:
                    time.sleep(2 ** attempt)
        
        return last_result