from typing import Union
from fairseq import tasks
from fairseq.data.dictionary import Dictionary
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel

import logging
import argparse
import torch
import torch.nn as nn
import pickle
import soundfile as sf
import torch.nn.functional as F
import yaml
import os, sys
import numpy as np


# from decoder_exps.decode_common import W2V2Decoder, Wav2VecCtc
# from wav2vecEncoder import Wav2VecCtc as CustomWav2VecCtc
# from fairseq.models.wav2vec.wav2vec2_asr import Wav2VecCtc

from inference import inference, inference_file, inference_online
# from inference import inference_file

logging.basicConfig()
logging.root.setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Recording(BaseModel):
    filename: str
    content_type: str
    content: list

class AudioClip(BaseModel):
    device: str
    time: str
    recording: Recording


# YAML_FILE = "config/base_org.yaml"

# args = dict()
# os.path.abspath(os.path.dirname(__file__))
# with open(YAML_FILE, 'r') as f:
#     args.update(yaml.safe_load(f))


###############################################################################
## FastAPI
###############################################################################

app = FastAPI()

@app.get("/")
def root():
    return {"message" : "Hello World!!!"}

@app.post("/test")
def post_test(audioClip: AudioClip):
    output = ""
    print("input audio? = ", type(audioClip.recording.content))
    audio = np.array(audioClip.recording.content).squeeze()

    print("in test func, audio = ", type(audio), audio.shape)

    feats = get_feature(audio)
    print("in test section, feats = ", type(feats), feats.shape)
    output = inference(feats)


    return {"output" : output}


@app.post("/inference")
def post_inference(audioClip: AudioClip):

    output = inference(audioClip.recording.content)

    return {"output" : output}

@app.post("/online")
def post_inference(audioClip: AudioClip):

    output = inference_online(audioClip.recording.content)

    return {"output" : output}


@app.post("/inference_file")
# def inference_file(file: UploadFile = File(...)):
def post_inference_file():
    # data = pickle.load(file)

    '''
    control formatting
    if data.format != 'wav':
        do_formatting()
    '''

    ## run model
    print('in Inference Start')
    output = ''
    output = inference_file()
    
    return {"output" : output}



if __name__ == '__main__':
    print('this is main')
    print(inference(args["wav_path"]))