#!/usr/bin/env python3
"""
Mashup Detection System
Identifies potential mashups in music metadata by analyzing delimiters,
validating component songs against a reference database, and scoring
based on multiple criteria.

Memory-optimized for large reference datasets.
"""

import json
import re
import gc
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field, asdict
import argparse
from difflib import SequenceMatcher
import ast
from tqdm import tqdm


@dataclass
class ComponentData:
    """Represents an extracted song component from a potential mashup"""
    raw_text: str
    song_name: str
    metadata: List[str]
    db_match: Optional[Dict] = None
    match_type: Optional[str] = None  # 'exact', 'fuzzy_XX', None


@dataclass
class ExtractedData:
    """Contains all extracted information from title parsing"""
    delimiter_type: Optional[str] = None
    delimiter_count: int = 0
    components: List[ComponentData] = field(default_factory=list)


@dataclass
class Scores:
    """Scoring categories for mashup detection"""
    delimiter_found: Optional[str] = None      # delimiter string found
    songs_validated: int = 0                    # 0-3: how many songs found in DB
    mashup_keyword: Optional[str] = None        # keyword string found
    known_mashup_source: Optional[str] = None   # source name (artist/label)


@dataclass
class AnalysisResult:
    """Complete analysis result for a single entry"""
    original_entry: Dict
    scores: Scores
    extracted_data: ExtractedData
    detection_notes: List[str]


class MashupDetector:
    """Main class for detecting mashups in music metadata"""
    
    # Delimiters in order of confidence
    DELIMITERS = [
        (' vs. ', 'vs.'),
        (' vs ', 'vs'),
        (' Vs. ', 'vs.'),
        (' Vs ', 'vs'),
        (' VS ', 'vs'),
        (' x ', 'x'),
        (' X ', 'x'),
        (' / ', '/'),
        (' // ', '//'),
        (' + ', '+'),
        (' versus ', 'versus'),
    ]
    
    # Known mashup artists
    KNOWN_MASHUP_ARTISTS = {
        'girl talk', 'gregg gillis', 'danger mouse', '2manydjs', '2 many djs',
        'party ben', 'the kleptones', 'eric kleptone', 'go home productions',
        'mark vidler', 'ghp', 'dj morgoth', 'madeon', 'super mash bros',
        'white panda', 'the white panda', 'kap slap', 'jared lucas',
        'mashd n kutcher', 'norwegian recycling', 'peter bull', 'djbc',
        'dean gray', 'team9', 'wax audio', 'tom compagnoni', 'electrosound',
        'freelance hellraiser', 'bootie mashup', 'adriana a', 'mashup germany',
        'the hood internet', 'psycosis', 'neutral bling hotel', 'cecil otter',
        'swiss andy', 'wugazi', 'dj chetas', 'laidback luke', 'steve angello',
        'dj vimto', 'fragma'
    }
    
    # Known mashup/bootleg labels
    KNOWN_MASHUP_LABELS = {
        'bootie mashup', 'jacked recordings', 'formation months series',
        'mashstix', 'dj pool records'
    }
    
    # Mashup keywords by strength
    MASHUP_KEYWORDS = {
        'strong': ['mashup', 'mash-up', 'mash up', 'bootleg', 'white label'],
        'medium': ['blend', 'edit', 'rework', 'bastard pop'],
        'weak': ['remix', 'mix']
    }
    
    def __init__(self, reference_path: Optional[str] = None):
        """Initialize detector with optional reference database"""
        self.has_reference = reference_path is not None
        if self.has_reference:
            self.reference_db = self._load_reference(reference_path)
            self._build_search_index()
        else:
            self.reference_db = []
            self.title_index = {}
            self.title_normalized_index = {}
    
    def _normalize_entry(self, entry: Dict) -> Dict:
        """Normalize entry to standard format, handling multiple input formats
        
        Supports multiple formats:
        1. Format 1: {'title': '...', 'artists': [...], 'label': '...'}
           - artists can be list of strings OR list of dicts with 'name' key
           - Example: {'title': 'Song', 'artists': [{'name': 'Artist', 'id': '123'}]}
        2. Format 2: {'song_name': '...', 'version': '...', 'artist_names': [...], 'label': '...'}
        3. Format 3: {'title': '...', 'artist': '...' or [...], 'label': '...'}
        4. Format 4: {'title': '...', 'artist': '...', 'genre': '...', 'album': '...'}
        
        Field aliases:
        - genre/sos_genre: genre information
        - album/release_name: album/release information
        - version: appended to title if not already present
        
        Returns: Normalized entry with 'title', 'artist', 'label', 'genre', 'subgenre', 'album' fields
        """
        normalized = {}
        
        # Detect format and extract title
        if 'title' in entry:
            # Format 1, 3, 4
            normalized['title'] = entry['title']
            # If version field exists and not already in title, append it
            version = entry.get('version', '')
            if version and f"({version})" not in normalized['title']:
                normalized['title'] = f"{normalized['title']} ({version})"
        elif 'song_name' in entry:
            # Format 2
            song_name = entry['song_name']
            version = entry.get('version', '')
            if version:
                normalized['title'] = f"{song_name} ({version})"
            else:
                normalized['title'] = song_name
        else:
            normalized['title'] = ''
        
        # Extract artists
        if 'artist_names' in entry:
            # Format 2
            normalized['artist'] = entry['artist_names']
        elif 'artists' in entry:
            # Format 1
            normalized['artist'] = entry['artists']
        elif 'artist' in entry:
            # Format 3, 4
            normalized['artist'] = entry['artist']
        else:
            normalized['artist'] = []
        
        # Copy other common fields with aliases
        normalized['label'] = entry.get('label', '')
        
        # Genre: check both 'genre' and 'sos_genre'
        normalized['genre'] = entry.get('genre') or entry.get('sos_genre', '')
        normalized['subgenre'] = entry.get('subgenre', '')
        
        # Album: check both 'album' and 'release_name'
        normalized['album'] = entry.get('album') or entry.get('release_name', '')
        
        # Keep original entry data for reference
        normalized['_original'] = entry
        
        return normalized
    
    def _load_reference(self, path: str) -> List[Dict]:
        """Load reference database from JSONL file with memory optimization"""
        reference = []
        print(f"Loading reference database from {path}...")
        
        # First, count total lines for progress bar
        total_lines = self._count_lines_in_file(path)
        
        with open(path, 'r', encoding='utf-8') as f:
            with tqdm(total=total_lines, desc="Loading reference", unit="entries") as pbar:
                for line in f:
                    pbar.update(1)
                    
                    entry = self._parse_jsonl_line(line)
                    if not entry:
                        continue
                    
                    # Normalize entry to handle different formats
                    normalized = self._normalize_entry(entry)
                    
                    # Extract only necessary fields immediately
                    title = normalized.get('title', '').strip()
                    if not title:
                        continue
                    
                    # Extract artist names using the existing helper method
                    artist_names = self._extract_artists(normalized)
                    
                    # Create minimal entry with only necessary fields
                    minimal_entry = {
                        'title': title,
                        'artist': artist_names,
                    }
                    
                    reference.append(minimal_entry)
                    
                    # Explicitly delete the original entry to free memory
                    del entry
                    del normalized
        
        print(f"Loaded {len(reference)} valid entries from {total_lines} total lines")
        
        # Force garbage collection after loading
        gc.collect()
        
        return reference
    
    def _build_search_index(self):
        """Build search indices for faster lookups with memory optimization"""
        print("Building search indices...")
        self.title_index = {}
        self.title_normalized_index = {}  # For more flexible matching
        
        with tqdm(total=len(self.reference_db), desc="Building indices", unit="entries") as pbar:
            for idx, entry in enumerate(self.reference_db):
                pbar.update(1)
                    
                title = entry.get('title', '')
                title_lower = title.lower().strip()
                
                # Also create normalized version (remove common words, punctuation)
                title_normalized = self._normalize_title(title_lower)
                
                # Index by exact title - store index instead of full entry
                if title_lower not in self.title_index:
                    self.title_index[title_lower] = []
                self.title_index[title_lower].append(idx)  # Store index, not entry
                
                # Index by normalized title
                if title_normalized not in self.title_normalized_index:
                    self.title_normalized_index[title_normalized] = []
                self.title_normalized_index[title_normalized].append(idx)
        
        print(f"Built indices with {len(self.title_index)} unique titles")
        
        # Force garbage collection after indexing
        gc.collect()
    
    def _normalize_title(self, title: str) -> str:
        """Normalize title for flexible matching"""
        # Remove common words and punctuation
        stop_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for'}
        
        # Remove punctuation and split
        title_clean = re.sub(r'[^\w\s]', ' ', title.lower())
        words = title_clean.split()
        
        # Filter stop words and rejoin
        filtered_words = [w for w in words if w not in stop_words]
        return ' '.join(filtered_words)
    
    def _count_lines_in_file(self, path: str) -> int:
        """Count total lines in a file"""
        total_lines = 0
        with open(path, 'r', encoding='utf-8') as f:
            for _ in f:
                total_lines += 1
        return total_lines
    
    def _disambiguate_matches_by_artist(self, match_indices: List[int], hint_artist: str) -> Optional[Dict]:
        """Try to disambiguate multiple matches using artist hint
        Returns: Matched entry or None if no disambiguation possible
        """
        for idx in match_indices:
            match = self.reference_db[idx]
            match_artists = match.get('artist', [])
            if isinstance(match_artists, str):
                match_artists = [match_artists]
            
            for artist in match_artists:
                if hint_artist.lower() in artist.lower() or artist.lower() in hint_artist.lower():
                    return match
        return None
    
    def detect_delimiter(self, title: str) -> Tuple[Optional[str], int]:
        """
        Detect delimiter in title and return type and count
        Returns: (delimiter_type, count)
        """
        for delimiter_pattern, delimiter_name in self.DELIMITERS:
            count = title.count(delimiter_pattern)
            if count > 0:
                return delimiter_name, count
        return None, 0
    
    def split_by_delimiter(self, title: str, delimiter_type: str) -> List[str]:
        """Split title by the detected delimiter"""
        # Find the actual delimiter pattern for the type
        delimiter_pattern = None
        for pattern, name in self.DELIMITERS:
            if name == delimiter_type:
                delimiter_pattern = pattern
                break
        
        if delimiter_pattern:
            # Split and clean each component
            components = [comp.strip() for comp in title.split(delimiter_pattern)]
            return components
        return [title]
    
    def extract_metadata(self, text: str) -> Tuple[str, List[str]]:
        """
        Extract song name and metadata from a text component
        Returns: (song_name, metadata_list)
        """
        song_name = text
        metadata = []
        
        # Extract all parenthetical content
        paren_pattern = r'\([^)]+\)'
        paren_matches = re.findall(paren_pattern, text)
        
        # Remove parenthetical content to get clean song name
        song_name = re.sub(paren_pattern, '', text).strip()
        
        # Process each parenthetical match
        for match in paren_matches:
            # Remove outer parentheses
            content = match[1:-1].strip()
            metadata.append(content)
        
        return song_name, metadata
    
    def find_song_in_db(self, song_name: str, hint_artist: Optional[str] = None) -> Tuple[Optional[Dict], str]:
        """
        Find a song in the reference database
        Returns: (matched_entry, match_type)
        """
        if not self.has_reference:
            return None, None
            
        song_lower = song_name.lower().strip()
        
        # Try exact match first
        if song_lower in self.title_index:
            match_indices = self.title_index[song_lower]
            if len(match_indices) == 1:
                return self.reference_db[match_indices[0]], 'exact'
            elif len(match_indices) > 1 and hint_artist:
                # Multiple matches, try to disambiguate with artist hint
                disambiguated = self._disambiguate_matches_by_artist(match_indices, hint_artist)
                if disambiguated:
                    return disambiguated, 'exact'
            
            # Return first match if no disambiguation possible
            return self.reference_db[match_indices[0]], 'exact'
        
        # Try normalized title match
        song_normalized = self._normalize_title(song_lower)
        if song_normalized in self.title_normalized_index:
            match_indices = self.title_normalized_index[song_normalized]
            if len(match_indices) == 1:
                return self.reference_db[match_indices[0]], 'normalized'
            elif hint_artist:
                disambiguated = self._disambiguate_matches_by_artist(match_indices, hint_artist)
                if disambiguated:
                    return disambiguated, 'normalized'
            return self.reference_db[match_indices[0]], 'normalized'
        
        # Try fuzzy matching (only on a subset for performance)
        best_match_idx = None
        best_score = 0
        
        # Limit fuzzy matching to avoid performance issues with large dataset
        sample_size = min(10000, len(self.title_index))
        sampled_titles = list(self.title_index.keys())[:sample_size]
        
        for ref_title in sampled_titles:
            score = SequenceMatcher(None, song_lower, ref_title).ratio()
            if score > best_score and score >= 0.85:  # 85% similarity threshold
                best_score = score
                best_match_idx = self.title_index[ref_title][0]
        
        if best_match_idx is not None:
            return self.reference_db[best_match_idx], f'fuzzy_{int(best_score * 100)}'
        
        return None, None
    
    def _parse_jsonl_line(self, line: str) -> Optional[Dict]:
        """Parse a JSONL line that could be JSON or Python dict format
        Returns: Parsed dict or None if parsing fails
        """
        line = line.strip()
        if not line:
            return None
        
        try:
            # Handle both dictionary string and JSON formats
            if line.startswith('{') and "'" in line[:50]:
                # Likely Python dict format, convert to JSON
                return ast.literal_eval(line)
            else:
                return json.loads(line)
        except (json.JSONDecodeError, SyntaxError, ValueError):
            return None
    
    def _extract_artists(self, entry: Dict) -> List[str]:
        """Extract and normalize artist names from entry
        Handles both 'artist' and 'artists' fields, and nested dict formats
        Returns: List of artist name strings
        """
        artists_data = entry.get('artist', entry.get('artists', []))
        
        if isinstance(artists_data, str):
            return [artists_data]
        elif isinstance(artists_data, list):
            normalized = []
            for artist in artists_data:
                if isinstance(artist, dict) and 'name' in artist:
                    normalized.append(artist['name'])
                elif isinstance(artist, str):
                    normalized.append(artist)
            return normalized
        return []
    
    def _field_to_lowercase(self, value) -> str:
        """Convert a field value to lowercase string, handling lists and None
        Returns: Lowercase string (empty string if None)
        """
        if value is None:
            return ''
        elif isinstance(value, str):
            return value.lower()
        elif isinstance(value, list):
            # Join list elements with spaces
            return ' '.join(str(item).lower() for item in value if item)
        else:
            return str(value).lower()
    
    def check_mashup_keywords(self, entry: Dict) -> Optional[str]:
        """Check for mashup keywords in entry fields, returns keyword found or None
        
        Checks: title, artist, label, genre, sos_genre, subgenre, album, release_name
        """
        text_to_check = []
        
        # Collect all text fields (using helper to handle lists/strings)
        text_to_check.append(self._field_to_lowercase(entry.get('title')))
        
        artists = self._extract_artists(entry)
        text_to_check.extend([a.lower() for a in artists])
        
        text_to_check.append(self._field_to_lowercase(entry.get('label')))
        
        # Genre fields: genre, sos_genre, subgenre
        text_to_check.append(self._field_to_lowercase(entry.get('genre')))
        text_to_check.append(self._field_to_lowercase(entry.get('sos_genre')))
        text_to_check.append(self._field_to_lowercase(entry.get('subgenre')))
        
        # Album/release fields: album, release_name
        text_to_check.append(self._field_to_lowercase(entry.get('album')))
        text_to_check.append(self._field_to_lowercase(entry.get('release_name')))
        
        combined_text = ' '.join(text_to_check)
        
        # Check keywords by strength (prioritize strong keywords)
        for strength in ['strong', 'medium', 'weak']:
            for keyword in self.MASHUP_KEYWORDS[strength]:
                if keyword in combined_text:
                    return keyword
        
        return None
    
    def check_known_sources(self, entry: Dict) -> Optional[str]:
        """Check if artist or label is known for mashups, returns source name or None"""
        # Check artists
        artists = self._extract_artists(entry)
        
        for artist in artists:
            if artist.lower().strip() in self.KNOWN_MASHUP_ARTISTS:
                return artist
        
        # Check label
        label = entry.get('label', '').lower().strip()
        if label in self.KNOWN_MASHUP_LABELS:
            return label
        
        return None
    
    def analyze_entry(self, entry: Dict) -> AnalysisResult:
        """Analyze a single entry for mashup detection"""
        # Normalize entry to handle different input formats
        normalized_entry = self._normalize_entry(entry)
        title = normalized_entry.get('title', '')
        
        # Initialize result (keep original entry for output)
        result = AnalysisResult(
            original_entry=entry,
            scores=Scores(),
            extracted_data=ExtractedData(),
            detection_notes=[]
        )
        
        # Detect delimiter
        delimiter_type, delimiter_count = self.detect_delimiter(title)
        
        if delimiter_type:
            result.extracted_data.delimiter_type = delimiter_type
            result.extracted_data.delimiter_count = delimiter_count
            result.scores.delimiter_found = delimiter_type
            result.detection_notes.append(f"Found delimiter: '{delimiter_type}' ({delimiter_count} times)")
            # Split by delimiter
            raw_components = self.split_by_delimiter(title, delimiter_type)
            
            # Process each component
            validated_count = 0
            for raw_text in raw_components:
                song_name, metadata = self.extract_metadata(raw_text)
                
                # Try to find in database only if reference is available
                db_match = None
                match_type = None
                
                if self.has_reference:
                    # Look for artist hints in metadata
                    hint_artist = None
                    for meta in metadata:
                        if any(x in meta.lower() for x in ['remix', 'edit', 'bootleg']):
                            # Extract artist name from "Artist Remix" pattern
                            parts = meta.split()
                            if len(parts) > 1:
                                hint_artist = parts[0]
                    
                    db_match, match_type = self.find_song_in_db(song_name, hint_artist)
                    
                    if db_match:
                        validated_count += 1
                        result.detection_notes.append(f"Found '{song_name}' in database ({match_type})")
                
                component = ComponentData(
                    raw_text=raw_text,
                    song_name=song_name,
                    metadata=metadata,
                    db_match=db_match,
                    match_type=match_type
                )
                
                result.extracted_data.components.append(component)
            
            # Calculate songs validated score (only if reference available)
            if self.has_reference and len(raw_components) > 0:
                if validated_count >= len(raw_components):
                    result.scores.songs_validated = 3
                elif validated_count >= len(raw_components) * 0.5:
                    result.scores.songs_validated = 2
                elif validated_count > 0:
                    result.scores.songs_validated = 1
        
        # Check for mashup keywords
        result.scores.mashup_keyword = self.check_mashup_keywords(normalized_entry)
        if result.scores.mashup_keyword is not None:
            result.detection_notes.append(f"Contains mashup keyword: '{result.scores.mashup_keyword}'")
        
        # Check known sources
        result.scores.known_mashup_source = self.check_known_sources(normalized_entry)
        if result.scores.known_mashup_source is not None:
            # Determine if it's an artist or label
            artists = self._extract_artists(normalized_entry)
            is_artist = any(result.scores.known_mashup_source.lower() == artist.lower() 
                          for artist in artists)
            
            if is_artist:
                result.detection_notes.append(f"Known mashup artist: {result.scores.known_mashup_source}")
            else:
                result.detection_notes.append(f"Known mashup label: {result.scores.known_mashup_source}")
        
        return result
    
    def process_file(self, input_path: str, output_path: str):
        """Process entire JSONL file and output results"""
        results = []
        
        # Count total lines first for progress bar
        total_lines = self._count_lines_in_file(input_path)
        
        print(f"Processing {total_lines} entries from input file...")
        
        with open(input_path, 'r', encoding='utf-8') as f:
            with tqdm(total=total_lines, desc="Detecting mashups", unit="songs") as pbar:
                for line in f:
                    pbar.update(1)
                    
                    entry = self._parse_jsonl_line(line)
                    if not entry:
                        continue
                    
                    try:
                        result = self.analyze_entry(entry)
                        results.append(result)
                        
                    except Exception as e:
                        pbar.write(f"Error processing entry: {e}")
        
        # Filter results to only include entries with positive scores
        filtered_results = [
            result for result in results
            if (result.scores.delimiter_found is not None or
                result.scores.mashup_keyword is not None or
                result.scores.known_mashup_source is not None)
        ]
        
        print(f"Found {len(filtered_results)} entries with positive mashup indicators (out of {len(results)} total)")
        
        # Write filtered results with progress bar
        print(f"Writing {len(filtered_results)} results to {output_path}...")
        with open(output_path, 'w', encoding='utf-8') as f:
            with tqdm(total=len(filtered_results), desc="Writing results", unit="entries") as pbar:
                for result in filtered_results:
                    # Convert to dictionary for JSON serialization
                    result_dict = {
                        'original_entry': result.original_entry,
                        'scores': asdict(result.scores),
                        'extracted_data': {
                            'delimiter_type': result.extracted_data.delimiter_type,
                            'delimiter_count': result.extracted_data.delimiter_count,
                            'components': [
                                {
                                    'raw_text': comp.raw_text,
                                    'song_name': comp.song_name,
                                    'metadata': comp.metadata,
                                    'db_match': comp.db_match,
                                    'match_type': comp.match_type
                                }
                                for comp in result.extracted_data.components
                            ]
                        },
                        'detection_notes': result.detection_notes
                    }
                    
                    f.write(json.dumps(result_dict) + '\n')
                    pbar.update(1)
        
        print(f"✓ Saved {len(filtered_results)} entries with positive mashup indicators")
        if not self.has_reference:
            print("Note: Running without reference database - songs validation scores unavailable")
        return filtered_results


def main():
    parser = argparse.ArgumentParser(description='Detect mashups in music metadata')
    parser.add_argument('input', help='Input JSONL file with metadata')
    parser.add_argument('output', help='Output JSONL file for results')
    parser.add_argument('--reference', help='Reference JSONL file with song database (optional)')
    
    args = parser.parse_args()
    
    # Initialize detector with or without reference
    detector = MashupDetector(args.reference)
    
    # Process file
    detector.process_file(args.input, args.output)


if __name__ == '__main__':
    main()