import json
import infinity
import infinity.index
from infinity.common import REMOTE_HOST
import pandas as pds

def main():
    try:
        infinity_obj = infinity.connect(REMOTE_HOST)
        db = infinity_obj.get_database("default")
        # Drop my_table if it already exists
        db.drop_table("my_table", if_exists=True)
        # Create a table named "my_table"
        table = db.create_table("my_table", {"num": "integer", "body": "varchar", "vec": "vector, 4, float"}, None)
        table.insert([{"num": 1, "body": "unnecessary and harmful", "vec": [1.0, 1.2, 0.8, 0.9]}])
        table.insert([{"num": 2, "body": "Office for Harmful Blooms", "vec": [4.0, 4.2, 4.3, 4.5]}])
        res = table.output(["num","body"]).knn("vec", [3.0, 2.8, 2.7, 3.1], "float", "ip", 2).to_pl()
        
        pds_df = pds.DataFrame(res)
        json_data = pds_df.to_json() 

        table_obj = db.get_table("my_table")
        #query_b = table_obj.query_builder()
        qb_result = table_obj.output(["num", "body"]).knn("vec",[3.0, 2.8, 2.7, 3.1], "float","ip",3).to_pl()
        #qb_result1 = table_obj.output(["num", "body"]).match("body","bloom","topn=1").to_pl()
        qb_result1 = table_obj.output(["num", "body"]).fusion('rrf').to_pl()
        print("------json-------") 
        print(json_data) 
        print("------tabular - querybuilder-------") 
        print(qb_result) 

    except Exception as e:
        print(str(e))        

    # Print the JSON data
    #print("------json-------") 
    #print(json_data) 
    #print("------tabular - querybuilder-------") 
    #print(qb_result)


if __name__ == '__main__':
    main()