from collections import defaultdict
from pathlib import Path
import json


writeDocumention = Path(__file__).parent / "tool_docs.json"
queryExamples = Path(__file__).parent / "queryExamples.json"
queryParams = Path(__file__).parent / "queryParams.json"
queryDescription = Path(__file__).parent / "queryDescription.json"
queryTables = Path(__file__).parent / "queryTables.json"
functionDocumention = Path(__file__).parent / "functionDocumention.json"
retrieveTables = Path(__file__).parent.parent.parent / "docs/dataset_tables.json"

endpoints = {"query_lineage": "extended_lineage_new", "query_policy_results": "policy_initiative","query_products": "asbom_filtered","query_risk": "product_risk_1.0", "query_vulnerabilities": "API Vulnerabilities", "query_findings": "Findings"}



def addExample(function: str, question: str, querystr):
    """
    The function can accept querystr as either a JSON object (dict) or a string. If using a JSON object, set format="json"; otherwise, the default is to treat querystr as a plain string.
    """
    try: 
        with open(queryExamples, "r+") as f:
            data = json.load(f)
    except:
        data = {}
    if function not in endpoints.keys(): raise "Not a valid function"
    try: 
        querystr = json.loads(querystr)
    except:
        pass
    if function not in data: data[function] = []
    for i in range(len(data[function])): 
        if data[function][i]["input"] == question: return 
    data[function].append({"input": question,"output": querystr})
    with open(queryExamples, "w") as f:
        json.dump(data, fp=f, indent=2)

def addDescription(function: str, description: str):
    with open(queryDescription, "r+") as f:
        data = json.load(f)
        data[function] = description
    with open(queryDescription, "w") as f:
        json.dump(data, fp=f, indent=2)

def updateTables():
    with open(retrieveTables, "r") as f:
        data = json.load(f)
        tables = {}
        for key, value in endpoints.items():
            del data[value]["md"]
            tables[key] = data[value]
    with open(queryTables, "w") as f:
        json.dump(tables, fp=f, indent=2)

def getTables(key: str):
    with open(retrieveTables, "r") as f:
        res = []
        data = json.load(f)
        for item in data[key]["columns"]:
            res.append((item["column_name"], item["label"]))
    return res

def updateToolDescription():
    with open(queryExamples, "r") as f: 
        examples = json.load(f)
    with open(queryParams, "r") as f: 
        params = json.load(f)
    with open(queryDescription, "r") as f: 
        descriptions = json.load(f)
    with open(queryTables, "r") as f: 
        tables = json.load(f)
    endpointDocumention = defaultdict(list)
    with open(writeDocumention, "w") as f:
        for endpoint in endpoints.keys():
            endpointDocumention[endpoint].append({"name": endpoint})
            endpointDocumention[endpoint].append({"description": descriptions[endpoint]})
            #endpointDocumention[endpoint].append({"params": params})
            endpointDocumention[endpoint].append({"examples": examples[endpoint]})
            endpointDocumention[endpoint].append({"ColumnTables": tables[endpoint]})
        json.dump(endpointDocumention, fp=f, indent=2)

def parseDoc(): 
    with open(functionDocumention, "r") as f:
        data = json.load(f)
        name, description, tableName = data["name"], data["description"], data["tableName"] 
        endpoints[name] = tableName
        for item in data["examples"].values():
            addExample(function=name, question=item["question"], querystr=item["querystr"])
        addDescription(name, description=description)
        updateTables()
    updateToolDescription()
    
if __name__ == "__main__":
    updateTables()