"""Python interface to the C++ NAM engine subprocess.

Spawns nam_engine as a subprocess and communicates via stdin/stdout
pipes for real-time audio processing. Falls back to PyTorch if the
C++ engine is unavailable.
"""

from __future__ import annotations

import json
import logging
import os
import subprocess
import time
from pathlib import Path
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

ENGINE_PATH = Path(__file__).parent / 'nam_engine'
DEFAULT_BLOCK_SIZE = 256


class NAMEngineProcess:
    """Manages the C++ nam_engine subprocess for a single model."""

    def __init__(self, model_path: str | Path, block_size: int = DEFAULT_BLOCK_SIZE):
        self._model_path = Path(model_path)
        self._block_size = block_size
        self._proc: Optional[subprocess.Popen] = None
        self._static: bool = False
        self._sample_rate: float = 48000.0
        self._timing_samples: list[float] = []
        self._loaded: bool = False

    def start(self) -> bool:
        """Launch the engine subprocess."""
        if not self._model_path.exists():
            logger.error('Model not found: %s', self._model_path)
            return False

        if not ENGINE_PATH.exists():
            logger.error('NAM engine binary not found: %s', ENGINE_PATH)
            return False

        try:
            self._proc = subprocess.Popen(
                [str(ENGINE_PATH), str(self._model_path), str(self._block_size)],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
        except Exception as e:
            logger.error('Failed to start NAM engine: %s', e)
            return False

        # Wait for the ready signal on stderr
        ready_line = self._read_stderr_line(timeout=5.0)
        if ready_line is None:
            logger.error('NAM engine failed to start (timeout waiting for ready)')
            self.stop()
            return False

        if 'FAILED' in ready_line:
            logger.error('NAM engine failed: %s', ready_line)
            self.stop()
            return False

        # Parse ready line: "NAM engine: loaded <path> (static=<N>, sr=<N>)"
        #                     "NAM engine: block_size=<N>, ready"
        if 'static=1' in ready_line:
            self._static = True
        logger.info('NAM engine ready: %s', ready_line.strip())
        
        # Read the block_size line too
        ready_line2 = self._read_stderr_line(timeout=2.0)
        if ready_line2:
            logger.info('NAM engine ready2: %s', ready_line2.strip())

        self._loaded = True
        return True

    def process(self, audio_block: np.ndarray) -> np.ndarray:
        """Process a block of audio through the NAM engine.

        Args:
            audio_block: float32 numpy array of shape (N,) or (1, N)

        Returns:
            float32 numpy array of same shape.
        """
        if self._proc is None or self._proc.stdin is None or self._proc.stdout is None:
            return audio_block  # passthrough

        # Ensure 1D
        was_2d = audio_block.ndim == 2
        if was_2d:
            audio_block = audio_block[0]

        # Ensure float32
        if audio_block.dtype != np.float32:
            audio_block = audio_block.astype(np.float32)

        start = time.perf_counter()

        # Write block to engine
        self._proc.stdin.write(audio_block.tobytes())
        self._proc.stdin.flush()

        # Read processed block
        raw = self._proc.stdout.read(audio_block.nbytes)
        if len(raw) != audio_block.nbytes:
            logger.warning('NAM engine short read: got %d bytes, expected %d', 
                len(raw), audio_block.nbytes)
            return audio_block  # passthrough on error

        out = np.frombuffer(raw, dtype=np.float32).copy()

        # Reshape back if input was 2D
        if was_2d:
            out = out[np.newaxis, :]

        # Track timing
        elapsed_ms = (time.perf_counter() - start) * 1000
        self._timing_samples.append(elapsed_ms)
        if len(self._timing_samples) > 200:
            self._timing_samples = self._timing_samples[-100:]

        return out

    def stop(self):
        """Terminate the engine subprocess."""
        if self._proc is not None:
            try:
                self._proc.terminate()
                self._proc.wait(timeout=3)
            except Exception:
                self._proc.kill()
                self._proc.wait(timeout=1)
            self._proc = None
        self._loaded = False

    def _read_stderr_line(self, timeout: float = 5.0) -> Optional[str]:
        """Read a line from stderr with timeout."""
        if self._proc is None or self._proc.stderr is None:
            return None
        # Poll until data available or timeout
        import select
        import sys
        
        # Use polling loop
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            # Check if process is still alive
            if self._proc.poll() is not None:
                # Read remaining stderr
                remaining = self._proc.stderr.read().decode('utf-8', errors='replace')
                logger.error('Engine exited with code %d: %s', self._proc.returncode, remaining)
                return 'FAILED: process exited'
            
            # Try non-blocking read
            line = self._proc.stderr.readline()
            if line:
                return line.decode('utf-8', errors='replace')
            
            time.sleep(0.05)  # 50ms poll interval
        
        return None  # timeout

    @property
    def is_loaded(self) -> bool:
        return self._loaded

    @property
    def is_static(self) -> bool:
        return self._static

    @property
    def avg_inference_ms(self) -> float:
        if not self._timing_samples:
            return 0.0
        return float(np.mean(self._timing_samples[-50:]))

    def __del__(self):
        self.stop()


# ── Simple test ────────────────────────────────────────────────────────

if __name__ == '__main__':
    import sys
    logging.basicConfig(level=logging.INFO)
    
    model = sys.argv[1] if len(sys.argv) > 1 else 'models/nam/clean.nam'
    block_size = int(sys.argv[2]) if len(sys.argv) > 2 else 256
    
    engine = NAMEngineProcess(model, block_size)
    if not engine.start():
        print('FAILED to start engine')
        sys.exit(1)
    
    print(f'Engine loaded: static={engine.is_static}')
    
    # Benchmark
    block = np.random.randn(block_size).astype(np.float32) * 0.1
    
    # Warmup
    for _ in range(10):
        engine.process(block)
    
    # Timed
    times = []
    for _ in range(500):
        t0 = time.perf_counter()
        engine.process(block)
        times.append((time.perf_counter() - t0) * 1000)
    
    print(f'Avg: {np.mean(times):.3f} ms  Max: {np.max(times):.3f} ms  Min: {np.min(times):.3f} ms')
    print(f'Engine reported avg: {engine.avg_inference_ms:.3f} ms')
    
    engine.stop()
