import dataiku
import base64
import pandas as pd


def write_question_answer_sql(client, dataset_name, connection, connection_type, project_key, question, answer):
    """
    Save data into a SQL like dataset
    Args:
        client: the dataiku client
        dataset_name: name of the SQL dataset
        connection: name of the SQL connection used for saving
        connection_type: type of connection
        project_key: project key
        question: the question
        answer: the answer
    """
    dataset = dataiku.Dataset(dataset_name)
    table_name = dataset.get_location_info().get('info', {}).get('table')
    value_string = f"('{base64.b64encode(question.encode('utf-8')).decode('utf-8')}', '{base64.b64encode(answer.encode('utf-8')).decode('utf-8')}')"
    sql = f"""INSERT INTO "{table_name}" (question, answer) VALUES {value_string}"""
    client.sql_query(sql, connection=connection, type=connection_type, project_key=dataset.project_key,
                     post_queries=['COMMIT'])


def write_question_answer_csv(dataset_name, question, answer):
    """
    Save data into a CSV like dataset
    Args:
        dataset_name: the CSV dataset
        question: the question
        answer: the answer
    """
    dataset = dataiku.Dataset(dataset_name)
    row = {
        "question": [f"""{question}"""],
        "answer": [f"""{answer}"""]
    }
    df = dataset.get_dataframe()
    df = df.append(pd.DataFrame(row))
    with dataset.get_writer() as writer:
        writer.write_dataframe(df)