import pandas as pd
import requests
from bs4 import BeautifulSoup
import re
import multiprocessing
import os # Import os to get process ID for debugging

def get_song_info_from_wikipedia(wikipedia_song_url):
    """
    Fetches track length from a Wikipedia song page.
    Args:
        wikipedia_song_url (str): The URL of the Wikipedia page for the song.
    Returns:
        tuple: (track_length, None) or (None, None) if not found.
    """
    # Ensure wikipedia_song_url is a string before proceeding
    if not isinstance(wikipedia_song_url, str) or not wikipedia_song_url.startswith('http'):
        return None, None

    try:
        response = requests.get(wikipedia_song_url, timeout=10)
        response.raise_for_status() # Raise an HTTPError for bad responses (4xx or 5xx)
    except requests.exceptions.RequestException as e:
        # print(f"Process {os.getpid()}: Error fetching {wikipedia_song_url}: {e}") # For debugging multiprocessing
        return None, None

    soup = BeautifulSoup(response.text, 'html.parser')

    track_length = None

    # --- Extract Track Length from Infobox ---
    # Common classes for infobox tables: 'infobox vcard', 'infobox', 'infobox vevent'
    infobox = soup.find('table', class_=re.compile(r'infobox'))
    if infobox:
        for row in infobox.find_all('tr'):
            header = row.find('th', scope='row')
            data = row.find('td')
            if header and data:
                header_text = header.get_text(strip=True).lower()
                if 'length' in header_text:
                    track_length = data.get_text(strip=True)
                    # Clean up common references like [1], [a] etc.
                    track_length = re.sub(r'\[.*?\]', '', track_length).strip()
                    break

    return track_length, None # Always return None for awards as we are no longer fetching them

def update_csv_with_wikipedia_info(input_csv_path, output_csv_path):
    """
    Reads the input CSV, fetches missing track_length from Wikipedia, and writes to output CSV.
    Args:
        input_csv_path (str): Path to the input CSV file.
        output_csv_path (str): Path to the output CSV file where updated data will be saved.
    """
    try:
        df = pd.read_csv(input_csv_path)
    except FileNotFoundError:
        print(f"Error: Input CSV file not found at '{input_csv_path}'")
        return
    except Exception as e:
        print(f"Error reading CSV file: {e}")
        return

    # Ensure 'track_length' column exists, add if missing
    if 'track_length' not in df.columns:
        df['track_length'] = None

    # Identify rows that need processing
    # Create a mask for rows where track_length is NaN AND wikipedia_song_url is not NaN
    needs_processing_mask = pd.isna(df['track_length']) & pd.notna(df['wikipedia_song_url'])
    songs_to_process_df = df[needs_processing_mask].copy()

    total_songs = len(df)
    songs_to_update_count = len(songs_to_process_df)
    processed_count = 0

    if songs_to_update_count == 0:
        print("No songs require track_length updates or no Wikipedia URLs found for missing lengths.")
        try:
            df.to_csv(output_csv_path, index=False)
            print(f"\nNo updates needed. CSV saved to '{output_csv_path}'")
        except Exception as e:
            print(f"Error writing CSV file when no updates needed: {e}")
        return

    # Detect number of CPU cores
    num_cores = multiprocessing.cpu_count()
    print(f"Detected {num_cores} CPU cores. Using a pool of {num_cores} processes.")

    # Prepare arguments for multiprocessing: list of Wikipedia URLs
    wikipedia_urls_for_processing = songs_to_process_df['wikipedia_song_url'].tolist()
    original_indices = songs_to_process_df.index.tolist() # Store original indices to update df

    # Use a multiprocessing Pool
    with multiprocessing.Pool(processes=num_cores) as pool:
        # imap_unordered allows us to process results as they come in, good for progress updates
        results_iterator = pool.imap_unordered(get_song_info_from_wikipedia, wikipedia_urls_for_processing)

        for i, (length, _) in enumerate(results_iterator):
            processed_count += 1
            original_index = original_indices[i] # Get the original index for the current result
            song_name = df.loc[original_index, 'track_name'] if 'track_name' in df.columns else 'N/A'
            artist_name = df.loc[original_index, 'artist'] if 'artist' in df.columns else 'N/A'

            print(f"Processing {processed_count}/{songs_to_update_count} (Total: {processed_count}/{total_songs}): {song_name} by {artist_name}")

            if length:
                df.at[original_index, 'track_length'] = length
                print(f"  - Updated track_length: {length}")
            else:
                print(f"  - Could not find track_length for {song_name}")

    # For songs that were not processed (track_length already present or no Wikipedia URL)
    # The original df remains untouched for these rows, which is the desired behavior.
    skipped_songs_count = total_songs - songs_to_update_count
    print(f"\nSkipped {skipped_songs_count} songs (track_length already present or no Wikipedia URL).")


    try:
        df.to_csv(output_csv_path, index=False)
        print(f"\nSuccessfully updated CSV and saved to '{output_csv_path}'")
    except Exception as e:
        print(f"Error writing updated CSV file: {e}")

# --- How to use the script ---
if __name__ == "__main__":
    # This block ensures that multiprocessing works correctly on Windows by protecting
    # the main execution block.
    multiprocessing.freeze_support()

    # IMPORTANT: Replace 'your_input_file.csv' with the actual name of your CSV file.
    # The script will create a new file named 'updated_songs.csv' with the results.
    input_csv = 'songs.csv'
    output_csv = 'updated_songs.csv'

    update_csv_with_wikipedia_info(input_csv, output_csv)

    print("\nScript finished. Check 'updated_songs.csv' for the results.")
    print("Remember to install the required libraries: pandas, requests, beautifulsoup4.")
    print("You can install them using pip: pip install pandas requests beautifulsoup4")