def loadDataSet(test_dir):
    f = open(test_dir,'r',encoding='utf-8')
    fname ='ts.csv'
    with open(fname,'w') as fine:
        for line in f.readlines():
            line1 =line.split()
            PROMPT=line1[3]
            linew = PROMPT+"\n"
            fine.write(linew)

def do_load(table_name, audio_dir,text_dir, model, milvus_client, mysql_cli):
    if not table_name:
        table_name = DEFAULT_TABLE
    vectors, names = extract_features(audio_dir, model)
    ids = milvus_client.insert(table_name, vectors)
    loadDataSet(text_dir)
    data = pd.read_csv("ts.csv")
    text = data['PROMPT'].tolist()
    milvus_client.create_index(table_name)
    mysql_cli.create_mysql_table(table_name)
    mysql_cli.load_data_to_mysql(table_name, format_data(ids, names,text))
    return len(ids)

def do_search(host,table_name, audio_path, model, milvus_client, mysql_cli):
    try:
        if not table_name:
            table_name = DEFAULT_TABLE
        feat = get_audio_embedding(audio_path)
        vectors = milvus_client.search_vectors(table_name, [feat], TOP_K)
        vids = [str(x.id) for x in vectors[0]]
        paths,text = mysql_cli.search_by_milvus_ids(vids, table_name)
        distances = [x.distance for x in vectors[0]]
        for i in range(len(paths)):
             tmp = "http://" + str(host) + "/data?audio_path=" + str(paths[i])
             paths[i] = tmp
        return vids, paths, distances,text
    except Exception as e:
        LOGGER.error(" Error with search : {}".format(e))
        sys.exit(1)

def create_mysql_table(self, table_name):
        sql = "create table if not exists " + table_name + "(milvus_id TEXT, audio_path TEXT,text TEXT) ENGINE=InnoDB DEFAULT CHARSET=utf8;"
        try:
            self.cursor.execute(sql)
            LOGGER.debug("MYSQL create table: {} with sql: {}".format(table_name, sql))
        except Exception as e:
            LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql))
            sys.exit(1)

def load_data_to_mysql(self, table_name, data):
        sql = "insert into " + table_name + " (milvus_id,audio_path,text) values (%s,%s,%s);"
        try:
            self.cursor.executemany(sql, data)
            self.conn.commit()
            LOGGER.debug("MYSQL loads data to table: {} successfully".format(table_name))
        except Exception as e:
            LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql))
            sys.exit(1)


def search_by_milvus_ids(self, ids, table_name):
        str_ids = str(ids).replace('[', '').replace(']', '')
        sql = "select * from " + table_name + " where milvus_id in (" + str_ids + ") order by field (milvus_id," + str_ids + ");"
        try:
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            results_path=[res[1] for res in results]
            results_text=[res[2] for res in results]
            LOGGER.debug("MYSQL search by milvus id.")
            return results_path,results_text
        except Exception as e:
            LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql))
            sys.exit(1)

class Item(BaseModel):
    Table: Optional[str] = None
    File:str
    Text:str

@app.post('/audio/load')
async def load_audios(item: Item):
    # Insert all the audio files under the file path to Milvus/MySQL
    try:
        total_num = do_load(item.Table, item.File,item.Text,MODEL, MILVUS_CLI, MYSQL_CLI)
        LOGGER.info("Successfully loaded data, total count: {}".format(total_num))
        return {'status': True, 'msg': "Successfully loaded data!"}
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400

@app.post('/audio/search')
async def search_audio(request: Request,Table: str = None, audio: UploadFile = File(...)):
    # Search the uploaded audio in Milvus/MySQL
    try:
        # Save the upload data to server.
        content = await audio.read()
        audio_path = os.path.join(UPLOAD_PATH, audio.filename)
        with open(audio_path, "wb+") as f:
            f.write(content)
        host = request.headers['host']
        ids, paths,text, distances= do_search(host,Table, audio_path, MODEL, MILVUS_CLI, MYSQL_CLI)
        names=[]
        names = text
        res = dict(zip(paths, zip(names, distances)))
        #res = sorted(res.items(), key=lambda item: item[1][1])
        LOGGER.info("Successfully searched similar audio!")
        return res
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400