prepare9D.py

import os
import json
import gzip
import re
import xml.etree.ElementTree as ET
from tqdm import tqdm
import multiprocessing
import time
import traceback
import logging

# Constants
OUTPUT_FOLDER = 'essential_discogs_data'
REJECTED_FOLDER = 'rejected_discogs_data' # Folder for rejected records raw xml
REJECTED_LOG_FILE = 'rejected_log.txt' # File to log rejection reasons
CHUNK_SIZE = 30000  # Number of records per output JSON file
# *** Modification: Reduced max rejected samples ***
MAX_REJECTED_SAMPLES = 100 # Max number of rejected records to save raw xml for
PROCESSING_POOL_SIZE = max(1, os.cpu_count() - 1) # Use most CPU cores, leave one free
IMAP_CHUNKSIZE = 250 # How many tasks to send to each worker process at a time

# Rejection Reason Constants
REASON_SUCCESS = None # Indicate success
REASON_MISSING_TITLE = "Missing or empty <title> element"
REASON_EMPTY_TITLE_CLEANED = "Title became empty after cleaning"
REASON_ZERO_TRACKS = "Rejected because no valid tracks were found/processed"
REASON_PARSE_ERROR = "Worker failed to parse XML fragment"
REASON_EXTRACTION_ERROR = "Unexpected error during data extraction"
REASON_WORKER_ERROR = "Unexpected error in worker process"


# --- Setup Logging ---
# Basic configuration to append to the log file
# We log INFO for explicit rejections and ERROR for exceptions
logging.basicConfig(filename=REJECTED_LOG_FILE,
                    level=logging.INFO,
                    format='%(asctime)s - ID: %(release_id)s - Reason: %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    force=True) # Use force=True if basicConfig might be called elsewhere implicitly


# --- Data Cleaning Functions ---
def clean_name(name):
    """Removes trailing '(number)' identifiers from artist/label names."""
    if name is None:
        return None
    return re.sub(r'\s*\(\d+\)$', '', name).strip()

def clean_title(title):
    """
    Cleans titles according to the revised logic:
    0. Strip whitespace.
    1. Reduce multiple consecutive parentheses to single ones.
    2. Removes outermost parentheses if the title is fully enclosed.
    3. Removes trailing parenthetical qualifiers (like '(Remastered)').
    """
    if title is None:
        return None
    processed_title = title.strip()
    for _ in range(5):
        original_len = len(processed_title)
        processed_title = processed_title.replace('((', '(').replace('))', ')')
        if len(processed_title) == original_len: break
    stripped_title = processed_title
    if stripped_title.startswith('(') and stripped_title.endswith(')'):
        processed_title = stripped_title[1:-1]
    else:
        processed_title = stripped_title
    final_title = re.sub(r'\s*\([^)]*\)$', '', processed_title)
    return final_title.strip()


# --- Core Data Extraction Logic ---
def extract_essential_data(record):
    """
    Extracts essential fields from a single Discogs <release> XML element.
    Returns a tuple: (data_dict_or_None, reason_constant_or_None)
    Logs rejection reasons internally, except for ZERO_TRACKS.
    """
    record_id = record.get('id', 'unknown')
    log_extra = {'release_id': record_id}

    try:
        # --- Title Processing ---
        title_elem = record.find('title')
        if title_elem is None or not title_elem.text:
            logging.info(REASON_MISSING_TITLE, extra=log_extra)
            return (None, REASON_MISSING_TITLE)
        title = clean_title(title_elem.text)
        if not title:
            logging.info(REASON_EMPTY_TITLE_CLEANED, extra=log_extra)
            return (None, REASON_EMPTY_TITLE_CLEANED)

        # --- Release Artists ---
        release_artists_elem = record.find('artists')
        release_artists = []
        if release_artists_elem is not None:
            for artist in release_artists_elem.findall('artist'):
                name_elem = artist.find('name')
                if name_elem is not None and name_elem.text:
                    cleaned = clean_name(name_elem.text.strip())
                    if cleaned: release_artists.append(cleaned)
        release_artists_lower = [a.lower() for a in release_artists] if release_artists else []
        is_various = release_artists_lower == ['various artists'] or release_artists_lower == ['various']

        # --- Other Fields ---
        genres_elem = record.find('genres')
        genres = []
        if genres_elem is not None:
            genres = [genre.text.strip() for genre in genres_elem.findall('genre') if genre is not None and genre.text]

        subgenres_elem = record.find('styles')
        subgenres = []
        if subgenres_elem is not None:
            subgenres = [style.text.strip() for style in subgenres_elem.findall('style') if style is not None and style.text]

        country_elem = record.find('country')
        country = country_elem.text.strip() if country_elem is not None and country_elem.text else None

        labels_elem = record.find('labels')
        labels = []
        if labels_elem is not None:
            for label in labels_elem.findall('label'):
                name = label.get('name')
                if name:
                    cleaned = clean_name(name.strip())
                    if cleaned: labels.append(cleaned)

        # --- Tracks Processing ---
        tracks = []
        tracklist = record.find('tracklist')
        if tracklist is not None:
            for track in tracklist.findall('track'):
                track_artists = []
                track_extraartists = []
                track_title = None
                duration = None

                track_title_elem = track.find('title')
                if track_title_elem is not None and track_title_elem.text:
                    track_title = clean_title(track_title_elem.text)
                    if not track_title: continue
                if track_title is None: continue

                track_artists_elem = track.find('artists')
                if track_artists_elem is not None:
                    for artist in track_artists_elem.findall('artist'):
                        name_elem = artist.find('name')
                        if name_elem is not None and name_elem.text:
                             cleaned = clean_name(name_elem.text.strip())
                             if cleaned: track_artists.append(cleaned)

                if not track_artists and release_artists and not is_various:
                    track_artists = list(release_artists)

                track_artists_lower = [a.lower() for a in track_artists] if track_artists else []
                if track_artists_lower == ['various artists'] or track_artists_lower == ['various']:
                    continue

                duration_elem = track.find('duration')
                duration = duration_elem.text.strip() if duration_elem is not None and duration_elem.text else None

                extraartists_elem = track.find('extraartists')
                if extraartists_elem is not None:
                    for artist in extraartists_elem.findall('artist'):
                        name_elem = artist.find('name')
                        if name_elem is not None and name_elem.text:
                             cleaned = clean_name(name_elem.text.strip())
                             if cleaned: track_extraartists.append(cleaned)

                tracks.append({
                    'title': track_title,
                    'artists': track_artists,
                    'duration': duration,
                    'extraartists': track_extraartists
                })

        # *** Apply Rule: Reject releases with 0 tracks found/processed ***
        # *** Modification: Do not log this specific rejection reason ***
        if not tracks:
             # logging.info(REASON_ZERO_TRACKS, extra=log_extra) # Removed logging call
             return (None, REASON_ZERO_TRACKS) # Still return None, but indicate reason

        # --- Release Date / Year ---
        released_elem = record.find('released')
        year = ''
        release_date = ''
        if released_elem is not None and released_elem.text:
            release_date_raw = released_elem.text.strip()
            if release_date_raw:
                 release_date = release_date_raw
                 if len(release_date_raw) >= 4 and release_date_raw[:4].isdigit():
                      potential_year = release_date_raw[:4]
                      if potential_year != '0000': year = potential_year
                      else: release_date = ''
                 elif not release_date_raw.replace('-', '').isdigit():
                     release_date = ''

        # --- Construct Final Data ---
        essential_data = {
            'release': title, 'year': year, 'release_date': release_date,
            'release_artists': release_artists, 'country': country, 'labels': labels,
            'tracks': tracks, 'genres': genres, 'subgenres': subgenres
        }
        return (essential_data, REASON_SUCCESS) # Return data and None reason for success

    except Exception as e:
        error_details = f"{REASON_EXTRACTION_ERROR}: {e.__class__.__name__}: {e}"
        logging.error(error_details, extra=log_extra) # Log exceptions as errors
        # traceback.print_exc()
        return (None, REASON_EXTRACTION_ERROR) # Return None and indicate reason

# --- Generator for XML Release Strings ---
def generate_release_strings(input_file):
    record_count = 0
    context = None
    try:
        with gzip.open(input_file, 'rt', encoding='utf-8') as f:
            context = ET.iterparse(f, events=('end',))
            for event, elem in context:
                if event == 'end' and elem.tag == 'release':
                    record_count += 1
                    try:
                        yield ET.tostring(elem, encoding='unicode', method='xml')
                    except Exception as to_string_e:
                         logging.warning(f"Could not convert element #{record_count} to string: {to_string_e}", extra={'release_id': elem.get('id', 'unknown')})
                    finally:
                         elem.clear()
            print(f"\nFinished reading XML stream. Total potential records considered: {record_count}")
    except ET.ParseError as e:
        print(f"\nXML Parse Error encountered near record #{record_count+1}: {e}")
        logging.critical(f"XML Parse Error near record #{record_count+1}: {e}", extra={'release_id': 'parse_error'})
    except gzip.BadGzipFile:
         print(f"\nError: Input file '{input_file}' is not a valid Gzip file.")
         logging.critical(f"Bad Gzip file: {input_file}", extra={'release_id': 'gzip_error'})
    except FileNotFoundError:
         print(f"\nError: Input file '{input_file}' not found.")
         logging.critical(f"Input file not found: {input_file}", extra={'release_id': 'filenotfound_error'})
    except StopIteration:
         print(f"\nWarning: XML file '{input_file}' might be empty or iteration stopped early (processed ~{record_count} records).")
    except Exception as e:
        print(f"\nAn unexpected error occurred while reading '{input_file}' after approx {record_count} records: {e}")
        logging.critical(f"Unexpected read error after {record_count} records: {e}", extra={'release_id': 'read_error'})
        traceback.print_exc()

# --- Worker Process Function ---
def worker_process_record(xml_string):
    """
    Function executed by worker processes. Parses XML and extracts data.
    Returns a tuple: (extracted_data_dict_or_None, reason_constant_or_None, original_xml_string)
    Logs specific errors internally.
    """
    record_id_str = 'unknown'
    reason = REASON_WORKER_ERROR # Default reason if early exit/unknown error
    data = None
    try:
        id_match = re.search(r'<release id="([^"]+)"', xml_string)
        if id_match: record_id_str = id_match.group(1)
        log_extra = {'release_id': record_id_str}

        record = ET.fromstring(xml_string)
        data, reason = extract_essential_data(record) # Gets data and reason
        return (data, reason, xml_string)
    except ET.ParseError as pe:
        reason = REASON_PARSE_ERROR
        logging.error(f"{reason}: {pe}", extra=log_extra)
        return (None, reason, xml_string)
    except Exception as e:
        reason = REASON_WORKER_ERROR
        logging.error(f"{reason}: {e.__class__.__name__}: {e}", extra=log_extra)
        return (None, reason, xml_string)

# --- Main Processing Orchestration ---
def process_large_xml_parallel(input_file, output_folder, rejected_folder, rejected_log, chunk_size):
    """
    Processes the large XML file in parallel. Logs rejection reasons (except ZERO_TRACKS).
    Saves raw XML for sampled rejected items (except ZERO_TRACKS).
    """
    if not os.path.exists(output_folder): os.makedirs(output_folder); print(f"Created output folder: {output_folder}")
    if not os.path.exists(rejected_folder): os.makedirs(rejected_folder); print(f"Created rejected records folder: {rejected_folder}")
    log_dir = os.path.dirname(rejected_log)
    if log_dir and not os.path.exists(log_dir): os.makedirs(log_dir)
    # Clear log file at start of run? Optional - using append mode by default.
    # with open(rejected_log, 'w') as f: f.write(f"Log started at {time.strftime('%Y-%m-%d %H:%M:%S')}...\n")
    print(f"Logging rejections (excluding zero tracks) to: {rejected_log}")

    essential_data_list = []
    chunk_index = 1
    total_processed_count = 0
    rejected_count_total = 0
    rejected_saved_count = 0 # Counts samples *actually saved*
    pool = None

    print(f"Starting parallel processing with {PROCESSING_POOL_SIZE} workers...")
    print(f"Processing file: {input_file}")
    print(f"Outputting valid data to: {output_folder}")
    print(f"Outputting sampled rejected records XML (excluding zero tracks) to: {rejected_folder}")
    print(f"Max rejected samples to save: {MAX_REJECTED_SAMPLES}")


    try:
        pool = multiprocessing.Pool(processes=PROCESSING_POOL_SIZE)
        results_iterator = pool.imap_unordered(
            worker_process_record,
            generate_release_strings(input_file),
            chunksize=IMAP_CHUNKSIZE
        )
        pbar = tqdm(desc="Processing records", unit=" records", mininterval=1.0, smoothing=0.1)

        for result_tuple in results_iterator:
            total_processed_count += 1
            pbar.update(1)

            if result_tuple is None:
                print("\nWarning: Received unexpected None result tuple from worker pool.")
                rejected_count_total += 1 # Count it as rejected, though reason unknown
                continue

            essential_data, reason, original_xml_string = result_tuple

            if essential_data is not None and reason == REASON_SUCCESS:
                # Successfully processed
                essential_data_list.append(essential_data)
                if len(essential_data_list) >= chunk_size:
                    output_file = os.path.join(output_folder, f'releases_{chunk_index}.json')
                    try:
                        with open(output_file, 'w', encoding='utf-8') as out_f:
                            json.dump(essential_data_list, out_f, indent=2)
                        essential_data_list = []
                        chunk_index += 1
                    except IOError as e: print(f"\nError writing chunk file {output_file}: {e}")
                    except TypeError as e: print(f"\nError serializing data for chunk {output_file}: {e}")

            else:
                # Record was rejected (essential_data is None)
                rejected_count_total += 1

                # *** Modification: Only save sample if reason is not ZERO_TRACKS ***
                if reason != REASON_ZERO_TRACKS:
                    # Logging is handled internally in extractor/worker now

                    # Save raw XML sample if limit not reached
                    if rejected_saved_count < MAX_REJECTED_SAMPLES:
                        rejected_saved_count += 1 # Increment only when saving
                        rejected_filename = os.path.join(rejected_folder, f'rejected_{rejected_saved_count}.xml')
                        try:
                            with open(rejected_filename, 'w', encoding='utf-8') as rej_f:
                                rej_f.write(original_xml_string)
                        except IOError as e:
                            print(f"\nError writing rejected file {rejected_filename}: {e}")
                            rejected_saved_count -= 1 # Decrement if save failed

        pbar.close()

    except KeyboardInterrupt:
        print("\nProcess interrupted by user. Terminating pool...")
        if pool: pool.terminate()
        print("Pool terminated.")
    except Exception as e:
         print(f"\nAn critical error occurred during parallel processing: {e}")
         traceback.print_exc()
         if pool: pool.terminate()
    finally:
        if pool: pool.close(); pool.join(); print("Processing pool closed.")

    # Write any remaining data
    if essential_data_list:
        output_file = os.path.join(output_folder, f'releases_{chunk_index}.json')
        try:
            with open(output_file, 'w', encoding='utf-8') as out_f:
                json.dump(essential_data_list, out_f, indent=2)
            print(f"\nWrote final chunk {chunk_index} ({len(essential_data_list)} records) to {output_file}")
        except IOError as e: print(f"\nError writing final chunk file {output_file}: {e}")
        except TypeError as e: print(f"\nError serializing data for final chunk {output_file}: {e}")

    print(f"\nFinished processing.")
    accepted_count = total_processed_count - rejected_count_total
    print(f"Total elements processed by workers: {total_processed_count}")
    print(f"Accepted records written: {accepted_count}")
    print(f"Rejected records encountered: {rejected_count_total}")
    # Note: rejected_saved_count only includes those saved *before* the limit was hit and *excluding* zero-track ones.
    print(f"Rejected records XML sampled and saved (excluding zero tracks): {rejected_saved_count} (max: {MAX_REJECTED_SAMPLES})")


# --- Main Execution ---
def main():
    input_file = None
    script_dir = os.path.dirname(os.path.abspath(__file__))
    print(f"Searching for Discogs file in: {script_dir}")
    for file in os.listdir(script_dir):
        if file.startswith('discogs_') and file.endswith('_releases.xml.gz'):
            input_file = os.path.join(script_dir, file); break
    if not input_file:
        print(f"Info: Discogs file not found in script directory. Searching current working directory: {os.getcwd()}")
        for file in os.listdir('.'):
             if file.startswith('discogs_') and file.endswith('_releases.xml.gz'):
                input_file = os.path.abspath(file); break
    if not input_file:
        print(f"Error: No valid Discogs releases XML file found.")
        return

    start_time = time.time()
    output_folder_path = os.path.join(script_dir, OUTPUT_FOLDER)
    rejected_folder_path = os.path.join(script_dir, REJECTED_FOLDER)
    rejected_log_path = os.path.join(script_dir, REJECTED_LOG_FILE)

    process_large_xml_parallel(input_file, output_folder_path, rejected_folder_path, rejected_log_path, CHUNK_SIZE)
    end_time = time.time()
    print(f"Total execution time: {end_time - start_time:.2f} seconds")

if __name__ == "__main__":
    multiprocessing.freeze_support()
    main()