data/msc_processor.py

import json
import os
from typing import Dict, List, Any
from .models import Conversation, Session

class MSCDataProcessor:
    """Process MSC dataset from local directory structure"""
    
    def __init__(self, root_dir: str):
        self.root_dir = root_dir
        self.dialogue_dir = os.path.join(root_dir, "msc_dialogue")
        self.personasummary_dir = os.path.join(root_dir, "msc_personasummary")
        
        print("📥 Loading MSC metadata...")
        self.init_personas_all = self._load_json(
            os.path.join(self.dialogue_dir, "init_persona_all.json")
        )
        self.session_summaries = self._load_json(
            os.path.join(self.dialogue_dir, "sessionlevel_summaries_subsample5.json")
        )
        self.summaries = self._load_json(
            os.path.join(self.dialogue_dir, "summaries_subsample5.json")
        )
        print("✅ Metadata loaded.")
    
    def _load_json(self, file_path: str) -> Dict:
        """Load JSON file with support for both standard JSON and JSON Lines"""
        if not os.path.exists(file_path):
            return {}
        
        with open(file_path, 'r', encoding='utf-8') as f:
            try:
                return json.load(f)
            except json.JSONDecodeError:
                # Handle JSON Lines format
                f.seek(0)
                data = {}
                for line in f:
                    if line.strip():
                        item = json.loads(line)
                        item_id = self._extract_item_id(item)
                        data[item_id] = item
                return data
    
    def _extract_item_id(self, item: Dict) -> str:
        """Extract item ID from various possible locations"""
        metadata = item.get('metadata', {})
        return (metadata.get('initial_data_id') or 
                item.get('initial_data_id') or 
                item.get('convai2_id') or 
                f"unknown_{id(item)}")
    
    def _format_dialog(self, dialog: List[Dict]) -> str:
        """Format dialog turns into structured text with normalized speaker roles"""
        formatted = []
        for idx, turn in enumerate(dialog):
            speaker = turn.get('id', 'Speaker')
            
            # Normalize speaker IDs to User/Assistant
            if speaker in ['bot_0', 'Speaker 1']:
                speaker = 'User'
            elif speaker in ['bot_1', 'Speaker 2']:
                speaker = 'Assistant'
            else:
                # Handle non-standard speaker IDs (e.g., numeric IDs)
                persona_text = turn.get('persona_text', '')
                if persona_text:
                    speaker = 'Assistant'
                else:
                    # Fallback: alternate based on turn position
                    speaker = 'User' if idx % 2 == 0 else 'Assistant'
            
            text = turn.get('text', '')
            formatted.append(f"{speaker}: {text}")
        
        return "\n".join(formatted)
    
    def _get_persona_summary(self, conv_id: str, session_id: int, persona_idx: int = 0) -> str:
        """Retrieve persona summary from sessionlevel_summaries_subsample5.json"""
        sess_key = str(session_id)
        if sess_key in self.session_summaries and conv_id in self.session_summaries[sess_key]:
            summaries = self.session_summaries[sess_key][conv_id]
            if persona_idx < len(summaries) and len(summaries[persona_idx]) > 0:
                return summaries[persona_idx][0]
        return ""
    
    def _get_persona_sentences(self, conv_id: str, session_id: int, persona_idx: int = 0) -> List[str]:
        """Retrieve individual persona sentences from summaries_subsample5.json"""
        sess_key = str(session_id)
        if sess_key in self.summaries and conv_id in self.summaries[sess_key]:
            summaries = self.summaries[sess_key][conv_id]
            if persona_idx < len(summaries):
                return summaries[persona_idx]
        return []
    
    def _get_persona_from_agg_list(self, dialog: List[Dict]) -> str:
        """Extract persona from aggregated persona list in dialog turns"""
        for turn in reversed(dialog):
            agg_list = turn.get('agg_persona_list', [])
            if agg_list:
                return ". ".join(agg_list)
        return ""
    
    def _get_init_personas(self, conv_id: str) -> List[List[str]]:
        """Retrieve initial personas from init_persona_all.json"""
        if conv_id in self.init_personas_all:
            raw = self.init_personas_all[conv_id]
            # Handle both dict and list formats
            if isinstance(raw, dict) and 'init_personas' in raw:
                return raw['init_personas']
            elif isinstance(raw, list):
                return raw
        return []
    
    def _validate_conversation(self, conv: Conversation) -> bool:
        """Validate conversation has sufficient data for training"""
        if len(conv.sessions) < 3:
            return False
        
        if not conv.sessions[-1].dialog_text:
            return False
        
        has_persona = any(s.persona_summary for s in conv.sessions)
        return has_persona
    
    def process_split(self, split: str = "train") -> List[Conversation]:
        """Process entire dataset for a given split and group sessions by conversation ID"""
        print(f"\n🔄 Processing {split} split...")
        
        conv_dict: Dict[str, Dict[int, Session]] = {}
        
        # Process msc_personasummary (sessions 1-4)
        for sess_id in range(1, 5):
            file_path = os.path.join(self.personasummary_dir, f"session_{sess_id}", f"{split}.txt")
            if not os.path.exists(file_path):
                continue
            
            print(f"  - Reading personasummary session {sess_id}...")
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if not line.strip():
                        continue
                    item = json.loads(line)
                    conv_id = item.get('initial_data_id')
                    if not conv_id:
                        continue
                    
                    if conv_id not in conv_dict:
                        conv_dict[conv_id] = {}
                    
                    dialog_text = self._format_dialog(item.get('dialog', []))
                    persona_summary = self._get_persona_summary(conv_id, sess_id, persona_idx=0)
                    persona_sentences = self._get_persona_sentences(conv_id, sess_id, persona_idx=0)
                    
                    session = Session(
                        session_id=sess_id,
                        dialog_text=dialog_text,
                        persona_summary=persona_summary,
                        persona_sentences=persona_sentences,
                        followup=item.get('followup'),
                        newfact=item.get('newfact'),
                        init_personas=None
                    )
                    conv_dict[conv_id][sess_id] = session
        
        # Process msc_dialogue (sessions 2-5)
        for sess_id in range(2, 6):
            file_path = os.path.join(self.dialogue_dir, f"session_{sess_id}", f"{split}.txt")
            if not os.path.exists(file_path):
                continue
            
            print(f"  - Reading dialogue session {sess_id}...")
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if not line.strip():
                        continue
                    item = json.loads(line)
                    conv_id = item.get('metadata', {}).get('initial_data_id')
                    if not conv_id:
                        continue
                    
                    if conv_id not in conv_dict:
                        conv_dict[conv_id] = {}
                    
                    dialog_text = self._format_dialog(item.get('dialog', []))
                    
                    # Prioritize session summaries, fallback to aggregated persona list
                    persona_summary = self._get_persona_summary(conv_id, sess_id, persona_idx=0)
                    if not persona_summary:
                        persona_summary = self._get_persona_from_agg_list(item.get('dialog', []))
                    
                    persona_sentences = self._get_persona_sentences(conv_id, sess_id, persona_idx=0)
                    
                    # Extract init_personas from nested structure
                    init_personachat = item.get('init_personachat', {})
                    if isinstance(init_personachat, dict):
                        init_personas = init_personachat.get('init_personas')
                    else:
                        init_personas = item.get('init_personas')
                    
                    session = Session(
                        session_id=sess_id,
                        dialog_text=dialog_text,
                        persona_summary=persona_summary,
                        persona_sentences=persona_sentences,
                        followup=None,
                        newfact=None,
                        init_personas=init_personas
                    )
                    conv_dict[conv_id][sess_id] = session
        
        # Convert to Conversation objects with validation
        conversations = []
        for conv_id, sessions_dict in conv_dict.items():
            if len(sessions_dict) < 3:
                continue
            
            sorted_sessions = [sessions_dict[i] for i in sorted(sessions_dict.keys())]
            init_personas = self._get_init_personas(conv_id)
            
            conv = Conversation(
                conv_id=conv_id,
                sessions=sorted_sessions,
                init_personas=init_personas
            )
            
            if self._validate_conversation(conv):
                conversations.append(conv)
        
        print(f"✅ Processed {len(conversations)} valid conversations from {split} split.")
        return conversations