import os
import glob
import json
import pandas as pd
import numpy as np

# ================= CONFIGURATION =================
DATASET_FOLDER = os.path.join(os.getcwd(), "dataset")
# =================================================

class Colors:
    HEADER = '\033[95m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

# Windows CMD color support fix
os.system('color')

def make_hashable(val):
    """Converts dialogue structures or lists into a stable string for hashing."""
    if val is None:
        return ""
    if isinstance(val, (list, np.ndarray, dict)):
        try:
            return json.dumps(val, sort_keys=True)
        except Exception:
            return str(val)
    return str(val).strip()

def load_file_to_df(filepath):
    """Loads various file formats and returns a unified DataFrame with row tracking."""
    ext = os.path.splitext(filepath)[1].lower()
    filename = os.path.basename(filepath)
    df = pd.DataFrame()
    
    try:
        if ext in ['.txt', '.md']:
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.read().split("\n\n")  # Split by double newline as rough sample block
            df = pd.DataFrame({'text': lines})
            df['row_index'] = df.index + 1
            
        elif ext == '.parquet':
            df = pd.read_parquet(filepath)
            df['row_index'] = df.index
            
        elif ext == '.csv':
            df = pd.read_csv(filepath)
            df['row_index'] = df.index + 2  # +2 accounts for header offset in CSV viewers
            
        elif ext in ['.json', '.jsonl']:
            try:
                df = pd.read_json(filepath, lines=True)
            except ValueError:
                df = pd.read_json(filepath)
            df['row_index'] = df.index
            
        if not df.empty:
            df['source_file'] = filename
            
    except Exception as e:
        print(f"{Colors.FAIL}Error reading {filename}: {e}{Colors.ENDC}")
        
    return df

def main():
    print(f"{Colors.HEADER}{'='*80}")
    print("             CROSS-FILE DATASET DUPLICATE TRACE SCANNER")
    print(f"{'='*80}{Colors.ENDC}")

    if not os.path.exists(DATASET_FOLDER):
        print(f"{Colors.FAIL}Error: Dataset folder '{DATASET_FOLDER}' not found.{Colors.ENDC}")
        return

    files = glob.glob(os.path.join(DATASET_FOLDER, "*.*"))
    if not files:
        print(f"{Colors.WARNING}No files found in '{DATASET_FOLDER}'.{Colors.ENDC}")
        return

    # 1. Load and aggregate all files
    all_data = []
    print("Loading and indexing raw dataset files...")
    for f in files:
        df = load_file_to_df(f)
        if not df.empty:
            print(f"  -> Indexed: {os.path.basename(f)} ({len(df)} rows)")
            all_data.append(df)

    if not all_data:
        print(f"{Colors.FAIL}No valid data could be parsed.{Colors.ENDC}")
        return

    # Combine everything into a global tracking DataFrame
    master_df = pd.concat(all_data, ignore_index=True)
    
    # Determine which column to audit
    target_column = None
    for col in ['conversations', 'messages', 'text']:
        if col in master_df.columns:
            target_column = col
            break
            
    if not target_column:
        # Fallback to the first non-meta column
        meta_cols = {'source_file', 'row_index'}
        valid_cols = [c for c in master_df.columns if c not in meta_cols]
        target_column = valid_cols[0] if valid_cols else master_df.columns[0]

    print(f"\nAuditing duplicates using key column: '{Colors.BOLD}{target_column}{Colors.ENDC}'")
    
    # Hash values for exact matching
    master_df['_hash'] = master_df[target_column].apply(make_hashable)
    
    # 2. Find duplicates (excluding empty rows)
    master_df = master_df[master_df['_hash'] != ""]
    
    # Group by the hash and find groups with size > 1
    dup_groups = master_df.groupby('_hash').filter(lambda x: len(x) > 1)
    
    if dup_groups.empty:
        print(f"\n{Colors.OKGREEN}🎉 Scan complete! No duplicate records found across any of the files.{Colors.ENDC}")
        return

    # 3. Report detailed trace
    unique_dupe_hashes = dup_groups['_hash'].unique()
    total_dupes_count = len(dup_groups) - len(unique_dupe_hashes)
    
    print(f"\n{Colors.WARNING}⚠️ Found {len(unique_dupe_hashes)} unique texts that have duplicated copies ({total_dupes_count} redundant entries total).{Colors.ENDC}")
    print("="*80)
    print("DETAILED DUPLICATE TRACE REPORT:")
    print("="*80)

    # Show up to 15 duplicate groups to keep console readable
    max_groups_to_show = 15
    for idx, h in enumerate(unique_dupe_hashes[:max_groups_to_show]):
        group = dup_groups[dup_groups['_hash'] == h]
        
        # Extract a clean text preview
        raw_val = group.iloc[0][target_column]
        if isinstance(raw_val, list):  # ShareGPT format preview
            preview = f"[ShareGPT Conversation with {len(raw_val)} turns]"
        else:
            preview = str(raw_val).replace('\n', ' ')[:80] + "..."

        print(f"\n{Colors.BOLD}Group #{idx+1} | Duplicate Content Preview:{Colors.ENDC}")
        print(f"  \"{Colors.OKGREEN}{preview}{Colors.ENDC}\"")
        print(f"  {Colors.BOLD}Found in following locations:{Colors.ENDC}")
        
        for _, row in group.iterrows():
            print(f"    - File: {Colors.WARNING}{row['source_file']}{Colors.ENDC} | Row/Index: {row['row_index']}")

    if len(unique_dupe_hashes) > max_groups_to_show:
        print(f"\n... and {len(unique_dupe_hashes) - max_groups_to_show} more duplicate groups.")

    print(f"\n{Colors.HEADER}{'='*80}{Colors.ENDC}")

if __name__ == "__main__":
    main()