import os
import sys
import time
import json
import shlex
import select
import logging
import threading
import subprocess as sp
import numpy as np
import tempfile

from ..utils import is_jetson, is_pi

_logger = logging.getLogger(__name__)
_logger.addHandler(logging.NullHandler())

_BAD_SIGS = (
    b"Impossible to convert between the formats",
    b"Failed to inject frame into filter network",
    b"Function not implemented",
    b"No device available for decoder",
    b"Invalid data found when processing input",
)

_ALLOWED_CODECS = {"h264", "hevc", "mpeg4", "mpeg2video", "vp8", "vp9"}


def _wait_proc(proc, timeout=5.0):
    """
    對長時間執行的 Popen 子行程：結束時必須顯式 wait() 才能避免 zombie。
    先 wait(timeout)，若逾時則 kill() 後再 wait(2)。
    """
    if proc is None:
        return
    try:
        proc.wait(timeout=timeout)
    except sp.TimeoutExpired:
        try:
            proc.kill()
            proc.wait(timeout=2.0)
        except Exception:
            pass
    except Exception:
        pass


# ---------------------------
# Tools for ffmpeg init
# ---------------------------
def has_decoder(name: str) -> bool:
    try:
        out = sp.check_output(["ffmpeg", "-hide_banner", "-decoders"], stderr=sp.STDOUT)
        return (name.encode() + b" ") in out or (b" " + name.encode() + b" ") in out
    except Exception:
        return False


def _run(cmd, timeout=6.0):
    """
    Conduct subprocess with stderr discarded to prevent noise interference in parsing.
    return (returncode:int, stdout:str)
    """
    env = os.environ.copy()
    # Avoid EGL error when no X on Jetson/container
    env.setdefault("EGL_PLATFORM", "surfaceless")
    p = sp.run(
        cmd,
        stdout=sp.PIPE,
        stderr=sp.DEVNULL,   # ignore stderr like 'No EGL Display'
        env=env,
        timeout=timeout,
        check=False,
    )
    out = p.stdout.decode("utf-8", "ignore")
    return p.returncode, out


def probe_stream_info(url: str, timeout: float = 6.0):
    """
    Get stream info from RTSP URL using ffprobe.
    Parse JSON output first; if fails, fallback to text parsing.
    Return (codec:str|None, width:int|None, height:int|None).
    這邊為None確認是否可以拉流
    
    Note: 如果連接失敗，會立即返回 (None, None, None)，避免快速重試導致 CPU 暴增。
    """

    # 對於 localhost，使用更短的超時，因為連接應該很快或立即失敗
    from urllib.parse import urlparse
    try:
        parsed = urlparse(url)
        host = parsed.hostname
        is_localhost = host in ('127.0.0.1', 'localhost', '::1', '0.0.0.0') if host else False
        if is_localhost:
            timeout = min(timeout, 2.0)  # localhost 使用更短s的超時
    except Exception:
        pass  # 如果解析失敗，使用原始 timeout

    # JSON parsing (most reliable)
    # 添加 -timeout 和 -rw_timeout 選項，避免 ffprobe 快速重試
    cmd = [
        "ffprobe", "-v", "quiet",
        "-rtsp_transport", "tcp", "-rtsp_flags", "prefer_tcp",
        "-timeout", str(int(timeout * 1000000)),  # 微秒單位，避免快速重試
        "-rw_timeout", str(int(timeout * 1000000)),  # 讀寫超時
        "-select_streams", "v:0",
        "-show_entries", "stream=codec_name,width,height",
        "-of", "json", url,
    ]
    rc, out = _run(cmd, timeout=timeout)
    if rc == 0 and out.strip():
        try:
            data = json.loads(out)
            streams = data.get("streams") or []
            if streams:
                s = streams[0]
                codec = (s.get("codec_name") or "").lower() or None
                w = s.get("width") or None
                h = s.get("height") or None
                if codec and codec in _ALLOWED_CODECS:
                    return codec, int(w) if w else None, int(h) if h else None
        except Exception:
            pass

    # Fallback to text parsing
    # This may be unreliable due to possible noise lines in output
    # Use nk=1:nw=1 to text output, get the last valid codec
    # 添加超時選項，避免快速重試
    cmd_fb = [
        "ffprobe", "-v", "error",
        "-rtsp_transport", "tcp", "-rtsp_flags", "prefer_tcp",
        "-timeout", str(int(timeout * 1000000)),  # 微秒單位
        "-rw_timeout", str(int(timeout * 1000000)),  # 讀寫超時
        "-select_streams", "v:0",
        "-show_entries", "stream=codec_name,width,height",
        "-of", "default=nk=1:nw=1", url,
    ]
    rc, out = _run(cmd_fb, timeout=timeout)
    if rc == 0 and out:
        lines = [x.strip().lower() for x in out.splitlines() if x.strip()]
        # Try to find codec first since it may have noise lines
        codec = next((x for x in lines if x in _ALLOWED_CODECS), None)
        
        def _to_int(s):
            try: return int(s)
            except: return None

        # Width, height may be in the last two lines (but double-check)
        w = _to_int(lines[-2]) if len(lines) >= 2 else None
        h = _to_int(lines[-1]) if len(lines) >= 1 else None
        if w is None or h is None:
            ints = [ _to_int(x) for x in lines ]
            ints = [i for i in ints if i is not None ]
            if len(ints) >= 2:
                w, h = ints[-2], ints[-1]
        return codec, w, h
    
    return None, None, None


# ---------------------------
# Tools for ffmpeg subprocess
# ---------------------------
def drain_stderr(proc):
    """
    Read and discard ffmpeg stderr to avoid blocking.
    Also print messages for debugging.
    """
    # Only drain but don't process/write to terminal unless necessary to save CPU.
    # FFmpeg is already called with -v error.
    for line in iter(proc.stderr.readline, b''):
        try:
            sys.stderr.write("[ffmpeg] " + line.decode("utf-8", "ignore"))
        except Exception:
            pass
    proc.stderr.close()  # 讀完後顯式關閉，避免 FD 延遲釋放


def build_cmd_multi_branch(rtsp_url, branches, fps=5, hwaccel=True, probe_timeout=5.0, branch_fifos=None, use_stdout_for_single=False):
    """
    Build FFmpeg command for multi-branch output mode.
    Uses filter_complex to split one stream into multiple branch outputs.
    
    Args:
        rtsp_url: RTSP URL
        branches: dict {branch_id: {WIDTH, HEIGHT, COLOR, LETTERBOX}}
        fps: Frame rate
        hwaccel: Use hardware acceleration
        probe_timeout: Timeout for codec probe
        branch_fifos: dict {branch_id: fifo_path} or None
        use_stdout_for_single: If True and only one branch, output to stdout instead of FIFO
    
    Returns:
        (cmd_list, codec)
    """
    fps = int(fps)
    
    # 1. Probe codec info
    codec, src_w, src_h = None, None, None
    try:
        codec, src_w, src_h = probe_stream_info(rtsp_url, timeout=probe_timeout)
        if codec:
            _logger.info(f"Stream probe: codec={codec}, resolution={src_w}x{src_h}")
    except Exception as e:
        _logger.warning(f"ffprobe failed: {e}")
    
    _is_pi = is_pi()
    codec_lower = (codec or "").lower()
    
    # 2. Setup Input and Decode Options
    input_opts = [
        "-rtsp_transport", "tcp",
        "-rtsp_flags", "prefer_tcp",
        "-use_wallclock_as_timestamps", "1",
    ]
    decode_opts = []
    
    if hwaccel:
        if _is_pi:
            # On RPi5, force DRM hardware acceleration if requested.
            # This ensures we don't fallback to software decoding.
            input_opts.extend(["-hwaccel", "drm", "-hwaccel_output_format", "drm_prime"])
        else:
            # Jetson HW Decoders
            nvmap = {
                "h264": "h264_nvv4l2dec", "avc": "h264_nvv4l2dec",
                "hevc": "hevc_nvv4l2dec", "h265": "hevc_nvv4l2dec",
                "vp8": "vp8_nvv4l2dec", "vp9": "vp9_nvv4l2dec"
            }
            nvdec = nvmap.get(codec_lower)
            if nvdec and has_decoder(nvdec):
                decode_opts = ["-c:v", nvdec]
    
    # RPi5 can do H.264 HW decode but it's often slower or more problematic than SW for small resolutions
    # so we focus on HEVC/H.265 here.

    # 3. Filter and Command Assembly
    num_b = len(branches)
    if num_b == 0: raise ValueError("No branches specified")

    # Helper for per-branch filter chain (scale, pad, format)
    def get_branch_filter(cfg, do_global, max_res, max_branch_w, max_branch_h):
        w, h, lb = int(cfg['WIDTH']), int(cfg['HEIGHT']), cfg.get('LETTERBOX')
        pix = 'bgr24' if cfg.get('COLOR', 'BGR').upper() == 'BGR' else 'rgb24'
        chain = []
        if lb:
            #do letterbox
            r = min(w / max_branch_w, h / max_branch_h)
            resize_w, resize_h = int(round(max_branch_w * r)), int(round(max_branch_h * r))
            chain.append(f"scale={resize_w}:{resize_h}:flags=fast_bilinear,pad={w}:{h}:(ow-iw)/2:(oh-ih)/2:color=black")
        elif not do_global or (w, h) != max_res:
            chain.append(f"scale={w}:{h}:flags=fast_bilinear")
        chain.append(f"format={pix}")
        return ",".join(chain), pix

    # Base Chain (FPS + HWDownload + optional Global Scale)
    max_branch_w = max(int(b["WIDTH"]) for b in branches.values())
    max_branch_h = max(int(b["HEIGHT"]) for b in branches.values())
    max_res = (max_branch_w, max_branch_h)

    base_parts = [f"fps={fps}"]
    if _is_pi and hwaccel and codec_lower in {"hevc", "h265"}:
        base_parts.append("hwdownload,format=nv12")
    
    do_global_downscale = False # Keep disabled as per current optimization
    if do_global_downscale:
        base_parts.append(f"scale={max_branch_w}:{max_branch_h}:flags=fast_bilinear")
    base_flt = ",".join(base_parts)

    # Common command components
    common = ["-v", "error", "-hide_banner", "-nostats"] + input_opts + decode_opts + \
             ["-i", rtsp_url, "-an", "-sn", "-dn", "-sws_flags", "fast_bilinear"]

    if num_b == 1 and use_stdout_for_single and not branch_fifos:
        # Simple Branch Output to stdout
        cfg = next(iter(branches.values()))
        chain, pix = get_branch_filter(cfg, do_global_downscale, max_res, max_branch_w, max_branch_h)
        cmd = ["ffmpeg"] + common + ["-vf", f"{base_flt},{chain}", "-pix_fmt", pix, "-vsync", "0", "-f", "rawvideo", "pipe:1"]
    else:
        # Multi-Branch Output via filter_complex
        flts = [f"[0:v]{base_flt},split={num_b}" + "".join(f"[branch{i}]" for i in range(num_b))]
        map_opts = []
        for i, (bid, cfg) in enumerate(sorted(branches.items())):
            chain, pix = get_branch_filter(cfg, do_global_downscale, max_res, max_branch_w, max_branch_h)
            flts.append(f"[branch{i}]{chain}[out{i}]")
            fifo = (branch_fifos or {}).get(bid)
            if not fifo: raise ValueError(f"FIFO path required for branch {bid}")
            map_opts.extend(["-map", f"[out{i}]", "-pix_fmt", pix, "-vsync", "0", "-f", "rawvideo", fifo])
            
        cmd = ["ffmpeg", "-y"] + common + ["-filter_complex", ";".join(flts)] + map_opts
    return cmd, codec



class FFmpegMultiBranchCapture:
    """
    Adapter for multi-branch output mode using single FFmpeg process.
    Uses filter_complex to split one RTSP stream into multiple branch outputs.
    Returns dict {branch_id: frame} on read_branches()
    
    CPU-optimized: Uses blocking IO with dedicated reader threads to avoid 
    select() polling overhead.
    """
    is_ffmpeg = True

    def __init__(
        self,
        branches,  # dict {branch_id: {WIDTH, HEIGHT, COLOR, LETTERBOX}}
        fps=5,
        hwaccel=True,
        fallback=False,
        first_frame_timeout=5.0,
        frame_timeout=3.0,
        zero_copy=True
    ):
        self.branches = branches
        self.fps = int(fps)
        self.hwaccel = bool(hwaccel)
        self.fallback = bool(fallback)
        self.first_frame_timeout = float(first_frame_timeout)
        self.frame_timeout = float(frame_timeout)
        self.zero_copy = bool(zero_copy)
        
        self._url = None
        self._proc = None
        self._opened = False
        self._reader_running = False
        
        # Frame sizes and buffers for each branch
        self._branch_buffers = {}
        self._branch_fifos = {}
        self._branch_queues = {}  # Queue for each branch to receive frames from reader threads
        self._reader_threads = {}  # Reader threads for each branch
        
        import queue as queue_module
        self._queue_module = queue_module
        
        for branch_id, branch_config in branches.items():
            w = int(branch_config['WIDTH'])
            h = int(branch_config['HEIGHT'])
            frame_size = w * h * 3
            self._branch_buffers[branch_id] = {
                'size': frame_size,
                'w': w,
                'h': h,
            }
            # Queue with maxsize=2 to avoid memory bloat while allowing some buffering
            self._branch_queues[branch_id] = queue_module.Queue(maxsize=2)

    def _create_fifos(self):
        """Create named pipes (FIFOs) for each branch output."""
        try:
            for branch_id in self.branches.keys():
                fifo = tempfile.NamedTemporaryFile(delete=False, suffix=f'.fifo.{branch_id}')
                fifo.close()
                os.unlink(fifo.name)
                os.mkfifo(fifo.name)
                self._branch_fifos[branch_id] = fifo.name
            return True
        except Exception as e:
            _logger.error(f"Failed to create FIFOs: {e}")
            self._cleanup_fifos()
            return False

    def _cleanup_fifos(self):
        """Clean up named pipes."""
        for branch_id, fifo_path in list(self._branch_fifos.items()):
            if fifo_path:
                try:
                    if os.path.exists(fifo_path):
                        os.unlink(fifo_path)
                except OSError as e:
                    _logger.debug("Cleanup FIFO %s: %s", fifo_path, e)
        self._branch_fifos.clear()

    def _start_proc(self, cmd):
        """Start FFmpeg process."""
        _logger.info("[cmd] %s", " ".join(shlex.quote(x) for x in cmd))
        p = sp.Popen(cmd, stdout=sp.DEVNULL, stderr=sp.PIPE, bufsize=10**8)
        threading.Thread(target=drain_stderr, args=(p,), daemon=True).start()
        return p

    def _fifo_reader_thread(self, branch_id, fifo_path):
        """Dedicated reader thread for a single FIFO."""
        buf_info = self._branch_buffers[branch_id]
        frame_size = buf_info['size']
        w, h = buf_info['w'], buf_info['h']
        q = self._branch_queues[branch_id]
        
        try:
            with open(fifo_path, 'rb') as f:
                while self._reader_running:
                    raw_data = f.read(frame_size)
                    if not raw_data:
                        break
                    if len(raw_data) < frame_size:
                        continue
                    
                    frame = np.frombuffer(raw_data, dtype=np.uint8).reshape((h, w, 3))
                    
                    try:
                        q.put_nowait(frame)
                    except self._queue_module.Full:
                        try:
                            q.get_nowait()
                            q.put_nowait(frame)
                        except:
                            pass
        except Exception as e:
            if self._reader_running:
                _logger.debug(f"FIFO reader thread for branch {branch_id} error: {e}")

    def _start_reader_threads(self):
        """Start reader threads for all branches."""
        self._reader_running = True
        for branch_id, fifo_path in self._branch_fifos.items():
            t = threading.Thread(
                target=self._fifo_reader_thread,
                args=(branch_id, fifo_path),
                daemon=True
            )
            t.start()
            self._reader_threads[branch_id] = t

    def _unblock_fifo_readers(self):
        """Open each FIFO's write end briefly to unblock reader threads stuck in open()."""
        for fifo_path in list(self._branch_fifos.values()):
            if fifo_path and os.path.exists(fifo_path):
                try:
                    # O_WRONLY|O_NONBLOCK: succeeds if there is a pending reader (a thread
                    # blocked in open()), which unblocks that reader's open() call.
                    fd = os.open(fifo_path, os.O_WRONLY | os.O_NONBLOCK)
                    os.close(fd)
                except OSError:
                    pass  # no pending reader, or FIFO already gone — OK

    def _stop_reader_threads(self):
        """Stop all reader threads."""
        self._reader_running = False
        for q in self._branch_queues.values():
            try:
                while not q.empty():
                    q.get_nowait()
            except:
                pass
        # Unblock reader threads that may be stuck in the blocking open(fifo, 'rb') call.
        # This happens when FFmpeg exits before it ever opens the FIFO write end
        # (e.g. RTSP auth failure), causing threads to wait in open() indefinitely.
        self._unblock_fifo_readers()

    def _wait_first_frames(self, timeout):
        """Wait for first frame from each branch queue."""
        deadline = time.time() + timeout
        first_frames = {}
        
        for branch_id, q in self._branch_queues.items():
            remaining = deadline - time.time()
            if remaining <= 0:
                return None
            try:
                frame = q.get(timeout=remaining)
                first_frames[branch_id] = frame
            except self._queue_module.Empty:
                return None
        
        return first_frames

    def open(self, rtsp_url):
        """Open RTSP stream and start multi-branch processing."""
        self.release()
        self._url = rtsp_url
        
        if not self._create_fifos():
            return False

        try:
            cmd, self.codec = build_cmd_multi_branch(
                rtsp_url, self.branches, self.fps, self.hwaccel,
                branch_fifos=self._branch_fifos
            )
            
            if cmd is None:
                return False
            
            self._start_reader_threads()
            
            self._proc = self._start_proc(cmd)
            time.sleep(0.2)
            
            if self._proc.poll() is not None:
                _logger.error(f"FFmpeg process exited with code {self._proc.returncode}")
                self._stop_reader_threads()
                self._proc = None
                return False
            
            first_frames = self._wait_first_frames(self.first_frame_timeout)
            
            if first_frames is None:
                if not (self.fallback and self.hwaccel):
                    try:
                        if self._proc and self._proc.poll() is None:
                            self._proc.terminate()
                    except Exception:
                        pass
                    _wait_proc(self._proc, timeout=5.0)
                    self._stop_reader_threads()
                    self._proc = None
                    return False

                _logger.info("Hardware decode failed, trying software decode...")
                try:
                    self._proc.kill()
                except Exception:
                    pass
                _wait_proc(self._proc, timeout=5.0)
                self._stop_reader_threads()

                # Join first-attempt threads before starting a second set.
                # Without this, _start_reader_threads() below would overwrite
                # _reader_threads entries, orphaning the first-attempt threads
                # with no way to track or join them later.
                for t in list(self._reader_threads.values()):
                    t.join(timeout=2.0)
                self._reader_threads.clear()

                cmd, self.codec = build_cmd_multi_branch(
                    rtsp_url, self.branches, self.fps, hwaccel=False,
                    branch_fifos=self._branch_fifos
                )

                if cmd is None:
                    return False

                self._start_reader_threads()
                self._proc = self._start_proc(cmd)
                time.sleep(0.2)
                
                if self._proc.poll() is not None:
                    _logger.error(f"FFmpeg process exited with code {self._proc.returncode}")
                    self._stop_reader_threads()
                    self._proc = None
                    return False
                
                first_frames = self._wait_first_frames(self.first_frame_timeout)
                
                if first_frames is None:
                    try:
                        self._proc.kill()
                    except Exception:
                        pass
                    _wait_proc(self._proc, timeout=5.0)
                    self._stop_reader_threads()
                    self._proc = None
                    return False

            self._opened = True
            return True
        finally:
            if not self._opened:
                self._stop_reader_threads()
                for t in list(self._reader_threads.values()):
                    t.join(timeout=2.0)
                self._reader_threads.clear()
                self._cleanup_fifos()

    def isOpened(self):
        """Check if capture is opened."""
        return bool(self._proc and self._proc.poll() is None and self._opened and self._reader_running)

    def get_branches(self):
        """Read frames from all branches."""
        if not self.isOpened():
            return False, None
        
        frames_dict = {}
        for branch_id, q in self._branch_queues.items():
            try:
                frame = q.get(timeout=self.frame_timeout)
                frames_dict[branch_id] = frame if self.zero_copy else frame.copy()
            except self._queue_module.Empty:
                return False, None
        
        return True, frames_dict

    def read(self):
        """Read single frame (for compatibility)."""
        ret, frames_dict = self.get_branches()
        if ret and frames_dict:
            first_branch_id = sorted(frames_dict.keys())[0]
            return True, frames_dict[first_branch_id]
        return False, None

    def release(self):
        """Release resources."""
        self._stop_reader_threads()

        if self._proc:
            try:
                self._proc.terminate()
            except Exception:
                pass
            _wait_proc(self._proc, timeout=5.0)
            self._proc = None

        # FFmpeg 已死，FIFO write end 已關，reader threads 的 blocking read 會解除
        # join 確保 threads 真正退出、FD 確實關閉後才繼續
        for t in list(self._reader_threads.values()):
            t.join(timeout=2.0)
        self._reader_threads.clear()

        self._cleanup_fifos()
        self._opened = False

    def get(self, prop):
        """Get property (for compatibility)."""
        try:
            import cv2
        except ImportError:
            return 0.0
        if self.branches:
            first_branch = next(iter(self.branches.values()))
            if prop == cv2.CAP_PROP_FRAME_WIDTH:  return float(first_branch['WIDTH'])
            if prop == cv2.CAP_PROP_FRAME_HEIGHT: return float(first_branch['HEIGHT'])
            if prop == cv2.CAP_PROP_FPS:          return float(self.fps)
            if prop == cv2.CAP_PROP_FOURCC:       return float(cv2.VideoWriter_fourcc(*self.codec)) if self.codec else 0.0
        return 0.0


