import json
import requests
from bs4 import BeautifulSoup
import time
from pathlib import Path
from urllib.parse import quote

def get_beatport_track_url(song_name, artist, data_source):
    """
    Scrape Beatport search results to find the track URL.
    Returns None if not found or if data_source is not beatport.
    """
    if data_source != 'beatport':
        return None
    
    # Clean up artist field (remove brackets)
    artist_clean = artist.strip('[]') if artist else ''
    
    # Build search query
    query = f"{song_name} {artist_clean}".strip()
    search_url = f"https://www.beatport.com/search?q={quote(query)}"
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
    }
    
    try:
        response = requests.get(search_url, headers=headers, timeout=15)
        response.raise_for_status()
        
        soup = BeautifulSoup(response.content, 'html.parser')
        
        # Look for track links in search results
        # Beatport uses <a> tags with class containing 'track' and href pattern /track/
        track_links = soup.find_all('a', href=True)
        
        for link in track_links:
            href = link.get('href', '')
            # Track URLs follow pattern: /track/track-name/track-id
            if '/track/' in href and href.count('/') >= 3:
                # Make sure it's a full URL
                if href.startswith('http'):
                    return href
                else:
                    return f"https://www.beatport.com{href}"
        
        return None
        
    except requests.exceptions.RequestException as e:
        print(f"Error fetching URL for '{song_name}': {e}")
        return None
    except Exception as e:
        print(f"Error parsing response for '{song_name}': {e}")
        return None


def process_jsonl(input_file, output_file, delay=2.0):
    """
    Read JSONL file, scrape Beatport URLs, and write to output JSONL.
    
    Args:
        input_file: Path to input JSONL file
        output_file: Path to output JSONL file
        delay: Delay between requests in seconds (default: 2.0)
    """
    input_path = Path(input_file)
    output_path = Path(output_file)
    
    if not input_path.exists():
        print(f"Error: Input file '{input_file}' not found!")
        return
    
    processed_count = 0
    found_count = 0
    
    print(f"Reading from: {input_file}")
    print(f"Writing to: {output_file}")
    print("-" * 60)
    
    with open(input_path, 'r', encoding='utf-8') as infile, \
         open(output_path, 'w', encoding='utf-8') as outfile:
        
        for line_num, line in enumerate(infile, 1):
            try:
                # Parse JSON line
                record = json.loads(line.strip())
                
                # Extract fields
                song_name = record.get('song_name', '')
                artist = record.get('artists', '')
                data_source = record.get('data_source', '')
                
                # Fetch URL
                webpage = get_beatport_track_url(song_name, artist, data_source)
                
                # Add webpage field to record
                record['webpage'] = webpage
                
                # Write to output file
                outfile.write(json.dumps(record, ensure_ascii=False) + '\n')
                
                processed_count += 1
                if webpage:
                    found_count += 1
                    print(f"✓ [{line_num}] Found: {song_name[:50]}...")
                    print(f"   URL: {webpage}")
                else:
                    print(f"✗ [{line_num}] Not found: {song_name[:50]}...")
                
                # Rate limiting - be respectful to the server
                if data_source == 'beatport':
                    time.sleep(delay)
                
            except json.JSONDecodeError as e:
                print(f"Warning: Skipping invalid JSON on line {line_num}: {e}")
                continue
            except Exception as e:
                print(f"Warning: Error processing line {line_num}: {e}")
                continue
    
    print("-" * 60)
    print(f"Processing complete!")
    print(f"Total records processed: {processed_count}")
    print(f"URLs found: {found_count}")
    print(f"URLs not found: {processed_count - found_count}")


if __name__ == "__main__":
    # Configuration
    INPUT_FILE = "/home/sara/task_data/cleaned_mashup_data_wout_ws.jsonl"
    OUTPUT_FILE = "/home/sara/task_data/cleaned_mashup_data_wout_ws_w_links.jsonl"
    DELAY_SECONDS = 2.0  # Delay between requests (be respectful!)
    
    process_jsonl(INPUT_FILE, OUTPUT_FILE, delay=DELAY_SECONDS)