from platform import processor

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

load_dotenv()






class SDKParams:
    def __init__(self, enhancement_level: float, sdk_key: str, model_id: str = "quail-vf-1.1-l-16khz"):
        self.enhancement_level = enhancement_level
        self.sdk_key = sdk_key
        self.model_path=  aic.Model.download(model_id, "./models")
        

def process_file_sdk(input_path: str, output_path: str, sdk_params: SDKParams):
    # Load audio
    audio, sr = librosa.load(input_path, sr=48000, mono=True)
    audio = audio.reshape(1, -1)  # Convert to planar format
    output = np.zeros_like(audio)
    load_dotenv()
    model = aic.Model.from_file(sdk_params.model_path)
    optimal_frames = model.get_optimal_num_frames(48000)
    config = aic.ProcessorConfig(sample_rate=sr, num_channels=1, num_frames= optimal_frames,allow_variable_frames=False )
    processor = aic.Processor(model, sdk_params.sdk_key, config)
    proc_ctx = processor.get_processor_context()
    #proc_ctx.set_parameter(aic.ProcessorParameter.EnhancementLevel, sdk_params.enhancement_level)
    #proc_ctx.set_parameter(aic.ProcessorParameter.VoiceGain, 1.5)
    #proc_ctx.set_parameter(aic.ProcessorParameter.Bypass, 0.0)
    chunk_size = optimal_frames
    for i in range(0, audio.shape[1], chunk_size):
        chunk = audio[:, i : i + chunk_size]
        # Pad last chunk if needed
        if chunk.shape[1] < chunk_size:
            last_chunk_size = chunk.shape[1]
            padded = np.zeros((1, chunk_size), dtype=audio.dtype)
            padded[:, : chunk.shape[1]] = chunk
            chunk = padded
            enhanced_chunk = processor.process(chunk)
            output[:, i : i + last_chunk_size] = enhanced_chunk[:, :last_chunk_size]
            break
        enhanced_chunk = processor.process(chunk)
        output[:, i : i + chunk_size] = enhanced_chunk[:, :chunk_size]
    # Save result
    sf.write(output_path, output.T, sr)