# sdk_audio.py (or keep inline)

from __future__ import annotations

import numpy as np
import librosa
import soundfile as sf
from dotenv import load_dotenv
import aic_sdk as aic
import os

from constants import MODEL_ID

load_dotenv()


class SDKWrapper:
    def __init__(self, model_id: str = MODEL_ID, models_dir: str = "./models"):
        if os.getenv("AIC_SDK_KEY") is None:
            raise RuntimeError("Missing AIC_SDK_KEY.")
        self.sdk_key = os.getenv("AIC_SDK_KEY")
        model_path = aic.Model.download(model_id, models_dir)
        self.model = aic.Model.from_file(model_path)

    def init_processor(self, sample_rate: int, enhancement_level: float, allow_variable_frames: bool = False, num_frames: int | None = None):
        self.processor_sample_rate = sample_rate
        self.processor_optimal_frames = self.model.get_optimal_num_frames(sample_rate)
        config = aic.ProcessorConfig(
            sample_rate=sample_rate,
            num_channels=1,
            num_frames=self.processor_optimal_frames if num_frames is None else num_frames,
            allow_variable_frames=allow_variable_frames,
        )
        processor = aic.Processor(self.model, self.sdk_key, config)
        processor.get_processor_context().set_parameter(
            aic.ProcessorParameter.EnhancementLevel, float(enhancement_level)
        )
        self.processor = processor

    def change_enhancement_level(self, enhancement_level: float):
        if not hasattr(self, "processor"):
            raise ValueError("Processor not initialized")
        self.processor.get_processor_context().set_parameter(
            aic.ProcessorParameter.EnhancementLevel, float(enhancement_level)
        )
    
    def process_sync(
        self,
        audio: np.ndarray,
    ) -> np.ndarray:
        """
            audio_array: 2D NumPy array with shape (num_channels, samples) containing audio data to be enhanced
        """
        if len(audio.shape) == 1:
            audio = audio.reshape(1, -1)
        if audio.shape[0] > 2 or len(audio.shape) != 2:
            raise ValueError("Expected audio with shape (n, frames)")
        out = np.zeros_like(audio)
        chunk_size = self.processor_optimal_frames
        n = audio.shape[1]
        for i in range(0, n, chunk_size):
            chunk = audio[:, i : i + chunk_size]
            if chunk.shape[1] < chunk_size:
                last = chunk.shape[1]
                padded = np.zeros((1, chunk_size), dtype=audio.dtype)
                padded[:, :last] = chunk
                enhanced = self.processor.process(padded)
                out[:, i : i + last] = enhanced[:, :last]
                break
            enhanced = self.processor.process(chunk)
            out[:, i : i + chunk_size] = enhanced[:, :chunk_size]
        return out
    
    def process_chunk(
        self,
        chunk: np.ndarray,
    ) -> np.ndarray:
        """
        Realtime processing: process a single chunk of audio and return enhanced chunk.
        """
        if not hasattr(self, "processor"):
            raise ValueError("Processor not initialized")
        chunk = np.asarray(chunk, dtype=np.float32).flatten()
        if chunk.size == 0:
            return chunk
        chunk_planar = chunk.reshape(1, -1)
        enhanced_planar = self.processor.process(chunk_planar)
        return enhanced_planar.flatten()