import gradio as gr
import pandas as pd
import json
import uuid
import requests

def convert_inputs(user_input, var):
    if var == "__random__":
        return [f"session_{uuid.uuid4()}" for _ in user_input]

    if not isinstance(var, list):
        var = [var]
        
    var += var*(len(user_input)//len(var) + 1)

    diff = len(var) - len(user_input) 
    var = var[:-diff]
    
    return var


def convert_output_vars(user_input, var):
    if isinstance(var, str):
        var = [[var]]
    elif not isinstance(var[0], list):
        var = [var]
        
    var += var*(len(user_input)//len(var) + 1)

    diff = len(var) - len(user_input) 
    var = var[:-diff]
    
    return var

def process_inputs(user_input_json, session_id_json, project_id, chat_url, update_vars_json, output_vars_json, token):
    # Convert JSON strings into Python lists
    user_input = json.loads(user_input_json)
    session_id = json.loads(session_id_json)
    # project_id = json.loads(project_id_json)
    update_vars = json.loads(update_vars_json)
    output_vars = json.loads(output_vars_json)

    session_id = convert_inputs(user_input,session_id)
    # project_id = convert_inputs(user_input,project_id)
    update_vars = convert_inputs(user_input,update_vars)
    output_vars = convert_output_vars(user_input,output_vars)

    unique_output = set()

    for x in output_vars:
        for y in x:
            unique_output.add(y)

    dict_output = {
        uni: [[] for _ in user_input] for uni in unique_output 
    }

    # --- Your function logic here ---
    df = pd.DataFrame({
        "user_input": user_input,
        "session_id": session_id,
        "project_id": project_id,
        "chat_url": [chat_url] * len(user_input),
        "update_variables": update_vars,
        "output_variables": output_vars,
        "token":token,
        } | dict_output
    )
        # "answer": [[] * len(x) if isinstance(x, list) else [] for x in user_input]
    # })  

    return df


def run_chat(chat_url, project_id, session_id, user_input, update_variables, output_variables, token):
    req = requests.post(
        chat_url,
        json={
            "project_id": project_id,
            "session_id": session_id,
            "user_input": user_input,
            "update_variables": update_variables,
            "output_variables": output_variables
        },
        headers={"Authorization" : f"Bearer {token}"}
    )

    print("[REQ CONTENT]", req.content)

    out = req.json()["data"]

    print("[OUT]",out.get("results"))

    return out
    #return {out_var:out[out_var] for out_var in output_variables}


def run_process(df):
    for index, row in df.iterrows():
        if isinstance(row["user_input"], list):
            answer = []
            for user_input in row["user_input"]:
                answer.append(
                    run_chat(
                        row["chat_url"],
                        row["project_id"],#: project_id,
                        row["session_id"],#: session_id,
                        user_input,#: user_input,
                        row["update_variables"],#: update_variables,
                        row["output_variables"],
                        row["token"]
                    )
                )
        else:
            answer = run_chat(
                row["chat_url"],
                row["project_id"],#: project_id,
                row["session_id"],#: session_id,
                row["user_input"],#: user_input,
                row["update_variables"],#: update_variables,
                row["output_variables"],
                row["token"]
            )

        if isinstance(answer, list):
            for ans in answer:
                for key,val in ans.items():
                    df.loc[index, key].append(val)
        else:
            for key,ans in answer.items():
                df.loc[index, key] = ans 
            
    csv_path = "output.csv"
    df.to_csv(csv_path, index=False)
    return csv_path, df


with gr.Blocks() as demo:
    gr.Markdown("## 💬 JSON Input ➜ DataFrame ➜ CSV Export")

    user_input = gr.Code(label="user_input (list[str] or list[list[str]])", language="json", value='["Hello", ["Hi", "How are you?"]]')
    session_id = gr.Code(label="""session_id (list[str] or str or "__random__")""", language="json", value='"__random__"')
    
    # project_id = gr.Code(label="project_id (list[str])", language="json", value='["p1", "p2"]')
    
    project_id = gr.Textbox(label="Project ID", value="[project_id]")
    chat_url = gr.Textbox(label="CHAT URL", value="https://example.com/chat")
    token = gr.Textbox(label="TOKEN", value="[token]")
    
    update_vars = gr.Code(label="update_variables (list[dict] or dict)", language="json", value='{}')
    output_vars = gr.Code(label="output_variables (list[list[str]] or list[str])", language="json", value='["results", "category_list"]')

    run_btn = gr.Button("Run Function")

    df_output = gr.Dataframe(label="Output DataFrame", interactive=True)
    
    process_btn = gr.Button("Process")
    
    file_output = gr.File(label="Download CSV")

    process_btn.click(
        fn=run_process,
        inputs=df_output,
        outputs=[file_output, df_output]
    )

    run_btn.click(
        fn=process_inputs,
        inputs=[user_input, session_id, project_id, chat_url, update_vars, output_vars, token],
        outputs=df_output
    )

demo.launch()