import json
from pprint import pprint


TASK_TREE_DEMO_FILE = "./utils/tasks_tree.json"


def fmt_task(tasks_tree: dict, task_id: str, history_log_mode: bool = False):
    task = tasks_tree.get(task_id)
    if not task:
        return None

    task_type = task.get("type")
    if task_type not in ("breakdown", "data_retrieval"):
        return None

    if task_type == "breakdown":
        if history_log_mode:
            return {
                "type": task_type,
                "task_id": task.get("task_id"),
                "sub_tasks": [i for i in [fmt_task(tasks_tree, task_id, history_log_mode) for task_id in task.get("sub_task_list", [])] if i],
                "status": task.get("status"),
            }
        else:
            return {
                "type": task_type,
                "clarified_task": task["clarified_task"],
                "description": task["description"],
                "assumptions": task.get("assumptions", []),
                "task_id": task.get("task_id"),
                "sub_tasks": [i for i in [fmt_task(tasks_tree, task_id, history_log_mode) for task_id in task.get("sub_task_list", [])] if i],
                "status": task.get("status"),
            }
    else:
        if history_log_mode:
            return {
                "type": task_type,
                "task_id": task.get("task_id"),
                "data": {
                    "columns": task.get("columns"),
                    "rows_count": len(task.get("rows", [])),
                },
                "chart_options": task.get("chart_options"),
                "sql": task.get("sql"),
                "sql_error": task.get("sql_error"),
                "status": task.get("status"),
            }
        else:
            return {
                "type": task_type,
                "clarified_task": task["clarified_task"],
                "description": task["description"],
                "assumptions": task.get("assumptions", []),
                "task_id": task.get("task_id"),
                "data": {
                    "columns": task.get("columns"),
                    "rows": task.get("rows"),
                },
                "chart_options": task.get("chart_options"),
                "sql": task.get("sql"),
                "sql_error": task.get("sql_error"),
                "status": task.get("status"),
            }


def fmt_tasks_tree(answer_id, tasks_tree):
    root_task = tasks_tree.get("0", {})

    result = {
        "agent": root_task.get("agent"),
        "status": root_task.get("status"),
        "answer_id": answer_id,
        "tasks": fmt_task(tasks_tree, "0"),
    }

    return result


if __name__ == "__main__":
    with open(TASK_TREE_DEMO_FILE) as f:
        tasks_tree = json.loads(f.read())["tasks_tree"]
        for k, v in tasks_tree.items():
            pprint(fmt_tasks_tree(k, v))