import argparse
import librosa
import sys
import numpy as np

from inference import predict_endpoint


def main():
    parser = argparse.ArgumentParser(description="Smart Turn v3.2 CPU inference")
    parser.add_argument("audio", help="Audio file to analyze")
    parser.add_argument("--start-ms", type=int, default=0, help="Current user turn start")
    parser.add_argument("--end-ms", type=int, help="Decision time (exclusive audio cutoff)")
    args = parser.parse_args()
    if args.start_ms < 0 or (args.end_ms is not None and args.end_ms <= args.start_ms):
        parser.error("Require 0 <= start-ms < end-ms")
    file_path = args.audio

    try:
        print(f"Loading audio file: {file_path}")
        # Load the audio file with original sample rate
        audio, sr = librosa.load(file_path, sr=None, mono=True)

        # Slice before resampling so no audio after the decision is included.
        start_sample = args.start_ms * sr // 1000
        end_sample = len(audio) if args.end_ms is None else args.end_ms * sr // 1000
        if start_sample >= len(audio) or end_sample > len(audio):
            raise ValueError("Requested time range lies outside the audio")
        audio = audio[start_sample:end_sample]

        print(f"Loaded audio with sample rate: {sr} Hz, duration: {len(audio) / sr:.2f} seconds")

        # If needed, resample to 16kHz (the model's expected sample rate)
        if sr != 16000:
            print(f"Resampling from {sr}Hz to 16000Hz")
            audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)

        # Convert audio to float32 if not already
        if audio.dtype != np.float32:
            audio = audio.astype(np.float32)

        # Make sure the audio is in the expected range [-1, 1]
        if np.max(np.abs(audio)) > 1.0:
            audio = audio / np.max(np.abs(audio))

        # Call the prediction function with both audio and sample rate
        print("Running endpoint prediction...")
        result = predict_endpoint(audio)

        # Display results
        print("\nResults:")
        print(f"Prediction: {'Complete' if result['prediction'] == 1 else 'Incomplete'}")
        print(f"Probability of complete: {result['probability']:.4f}")

    except Exception as e:
        print(f"Error processing audio file: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()