Gallery-Archivist/archivist/sites/furaffinity/management/commands/import_data.py

# your_app/management/commands/import_data.py

import os
import json
from blake3 import blake3
from tqdm.auto import tqdm

#from PIL import Image
from datetime import datetime

from django.core.management.base import BaseCommand
from django.utils.text import slugify
from django.core.exceptions import ObjectDoesNotExist
#from django.core.files import File
from django.utils import timezone

from sites.furaffinity.models import FA_Submission, FA_Submission_File, FA_Tags, FA_User, FA_Species, FA_Gender, FA_Mature

class Command(BaseCommand):
    help = 'Import data from JSON files in a folder or a single JSON file to the furaffinity archive'

    def add_arguments(self, parser):
        parser.add_argument('path', type=str, help='Path to the folder containing JSON files or a single JSON file')
        parser.add_argument('--delete', action='store_true', help='Delete imported files')


    def handle(self, *args, **kwargs):
        path = kwargs['path']
        delete = kwargs['delete']

        if os.path.isfile(path):
            self.process_json_file(path, delete)
        elif os.path.isdir(path):
            self.process_json_folder(path, delete)
        else:
            self.stdout.write(self.style.ERROR(f"The path '{path}' is not a valid file or folder."))
            return


    def process_json_file(self, file_path, delete):
        #self.stdout.write(self.style.NOTICE(f"Importing data from: {file_path}"))
        tqdm.write(self.style.NOTICE(f"Importing data from: {file_path}"))

        with open(file_path) as f:
            data = json.load(f)

        self.import_data(data, file_path, delete)

        #self.stdout.write(self.style.SUCCESS('Data imported successfully.'))
        tqdm.write(self.style.SUCCESS('Data imported successfully.'))
    
    def process_json_folder(self, folder_path, delete):
        if not os.path.exists(folder_path):
            #self.stdout.write(self.style.ERROR(f"The folder '{folder_path}' does not exist."))
            tqdm.write(self.style.ERROR(f"The folder '{folder_path}' does not exist."))
            return

        for root, dirs, files in tqdm(os.walk(folder_path), dynamic_ncols=True):
            for file_name in files:
                if file_name.endswith('.json'):
                    file_path = os.path.join(root, file_name)
                    self.process_json_file(file_path, delete)
                    


    def import_data(self, data, json_file_path, delete):    
        #self.stdout.write(self.style.NOTICE(data))
        
        submission, created = FA_Submission.objects.get_or_create(submission_id=data["id"])
        
        submission.media_url     = data["url"]
        submission.title         = data["title"]
        submission.description   = data["description"]
        
        submission.date = timezone.make_aware(datetime.strptime(data["date"], "%Y-%m-%d %H:%M:%S"))
        
        artist, created = FA_User.objects.get_or_create(artist_url=data["artist_url"], artist=data["artist"])
        submission.artist = artist
        
        species, created = FA_Species.objects.get_or_create(species=data["species"])
        submission.species = species
        
        for tag_name in data["tags"]:
            tag_slug = slugify(tag_name)
            try:
                # Check if the tag already exists in the database by name
                tag = FA_Tags.objects.get(tag_slug=tag_slug)
            except ObjectDoesNotExist:
                # If the tag does not exist, create a new tag and generate the slug
                tag = FA_Tags(tag=tag_name)
                tag.tag_slug = tag_slug

            tag.save()  # Save the tag (either new or existing)

            submission.tags.add(tag)  # Add the tag to the submission
        
        mature, created = FA_Mature.objects.get_or_create(mature=data["rating"])
        submission.mature_rating = mature
        
        submission.number_of_comments = data["comments"]
        submission.views = data["views"]
        
        gender, created = FA_Gender.objects.get_or_create(gender=data["gender"])
        submission.gender = gender
        
        submission.fa_theme = data["theme"]
        submission.fa_category = data["fa_category"]
        submission.gallery_type = data["subcategory"]
        submission.file_extension = data["extension"]
        submission.image_height = data["height"]
        submission.image_width = data["width"]
        
        
        file_path = json_file_path.removesuffix(".json")
        
        # Handle file import
        if os.path.exists(file_path):
            file_hash = self.compute_file_hash(file_path)
            try:
                file_instance = FA_Submission_File.objects.get(file_hash=file_hash)
                #self.stdout.write(self.style.NOTICE(f"Skipping: {file_path} file, already imported"))
                tqdm.write(self.style.NOTICE(f"Skipping: {file_path} file, already imported"))
            except FA_Submission_File.DoesNotExist:
                # If the file doesn't exist, create a new file instance and link it to the submission
                with open(file_path, 'rb') as file:
                    file_instance = FA_Submission_File()
                    file_instance.file_hash = file_hash
                    
                    file_name = os.path.basename(file_path)
                    Null, file_ext = os.path.splitext(file_name)
                    hash_file_name = file_hash + file_ext
                    file_instance.file.save(hash_file_name, file)
                    
                    file_instance.file_name = file_name
                    file_instance.save()
                    
                    tqdm.write(self.style.NOTICE(f"Import media file: {file_path}"))
            
            # Now link the image_instance to your_model_instance
            submission.file = file_instance
            
            
        else:
            #self.stdout.write(self.style.WARNING(f"File not found: {file_path}"))
            tqdm.write(self.style.WARNING(f"File not found: {file_path}"))   
                #file, created = FA_Submission_File.objects.get_or_create(file=file_hash)
                #print(file)
                #if not FA_Submission.objects.filter(image=image_hash).exists():
                #    with open(image_file_path, 'rb') as img_file:
                #        submission.image.save(os.path.basename(image_file_path), File(img_file), save=True)
                #    submission.image_hash = image_hash  # Save the image hash in the model
                #else:
                #    self.stdout.write(self.style.WARNING(f"Skipping duplicate image: {image_file_path}"))
                    
        #if os.path.exists(image_file_path):
        #    with open(image_file_path, 'rb') as img_file:
        #        submission.image.save(os.path.basename(image_file_path), File(img_file), save=True)
        
        submission.save()

        self.delete_imported_file(json_file_path, delete)
        self.delete_imported_file(file_path, delete)


    def compute_file_hash(self, file_path):
        try:
            # Compute BLAKE3 hash of the file
            hasher = blake3()
            with open(file_path, 'rb') as f:
                while chunk := f.read(65536):
                    hasher.update(chunk)
            return hasher.hexdigest()
        except Exception as e:
            #self.stdout.write(self.style.WARNING(f"Error computing file hash: {e}"))
            tqdm.write(self.style.WARNING(f"Error computing file hash: {e}"))
            return None


    # Delete the file if the --delete flag is used
    def delete_imported_file(self, file_path, delete):
        if delete:
            if os.path.exists(file_path):
                os.remove(file_path)
                tqdm.write(self.style.SUCCESS(f"Deleted: {file_path}"))
            else:
                tqdm.write(self.style.WARNING(f"File not found: {file_path}"))