#!/usr/bin/env python3
"""Print plain text version of ChatGPT's exported `conversations.json`

To get the JSON conversations, one option is to export your entire history via
https://help.openai.com/en/articles/7260999-how-do-i-export-my-chatgpt-history

Another option is to export a single message via tampermonkey
https://github.com/pionxzh/chatgpt-exporter
"""
import json
import sys

def read_conversations(file_path):
    with open(file_path, 'r') as file:
        data = json.load(file)
    # We are expecting a list of conversations.
    # But if we read a single conversation, let's return it as a list.
    if isinstance(data, dict):
        return [data]
    return data

def find_root_node(current_node, message_lut, verbose=False):
    root_node = None
    message_data = message_lut[current_node]
    while root_node is None:
        parent = message_data.get('parent')
        if parent is None:
            # We found the root node!
            root_node = message_data['id']
        else:
            # Move on to next message.
            message_data = message_lut[parent]
    if verbose:
        print(f"Root Message ID: {root_node}")
    return root_node

def print_message(message, depth):
    if message is None:
        return
    author_role = message['author']['role']
    if author_role not in {'user', 'assistant'}:
        return
    print("-" * 10)
    print(f"\n[{depth-1}] >>> {author_role.capitalize()}")
    content = message['content']
    parts = content['parts']
    # Question: Are there ever multiple parts? Let's find out.
    assert len(parts) == 1
    for part in parts:
        print(part)
    print()

def print_thread(message_id, message_lut, depth=0):
    message_data = message_lut[message_id]
    message = message_data.get('message')
    print_message(message, depth)
    for child_id in message_data['children']:
        print_thread(child_id, message_lut, depth=depth+1)

def print_conversation_header(conversation_data):
    print("=" * 80)
    print(f"Title: {conversation_data['title']}")
    print(f"URL: https://chat.openai.com/c/{conversation_data['id']}")

def print_conversations(conversations_data):
    for conversation_data in conversations_data:
        print_conversation_header(conversation_data)
        mapping = conversation_data['mapping']
        root_node = find_root_node(conversation_data['current_node'], mapping)
        print_thread(message_id=root_node, message_lut=mapping)

def main():
    file_path = 'conversations.json'
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
    conversations_data = read_conversations(file_path)
    print_conversations(conversations_data)

if __name__ == "__main__":
    main()