import json
import hashlib
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, List, Optional

class FinetuneDatasetCreator:
    """Creates JSONL finetune dataset from iteration XML files"""
    
    def __init__(
            self,
            iterations_dir: Path,
            system_prompt_file: Path,
            action_schema_file: Path,
            output_file: Path
        ):
        """
        Initialize the dataset creator
        
        Args:
            iterations_dir: Directory containing iteration XML files
            system_prompt_file: Path to system prompt file
            action_schema_file: Path to action schema file
            output_file: Path where JSONL dataset will be written
        """
        self.iterations_dir = Path(iterations_dir)
        self.system_prompt_file = Path(system_prompt_file)
        self.action_schema_file = Path(action_schema_file)
        self.output_file = Path(output_file)
        
        # Read and hash system prompt and action schema
        self.system_prompt = self.system_prompt_file.read_text()
        self.system_prompt_hash = self._calculate_hash(self.system_prompt)
        
        self.action_schema = self.action_schema_file.read_text()
        self.action_schema_hash = self._calculate_hash(self.action_schema)
        
    def _calculate_hash(self, content: str) -> str:
        """Calculate SHA256 hash of content"""
        return hashlib.sha256(content.encode()).hexdigest()
    
    def _parse_iteration_file(self, file_path: Path) -> Optional[Dict]:
        """
        Parse a single iteration XML file into a messages dictionary
        
        Returns None if hashes don't match or parsing fails
        """
        try:
            tree = ET.parse(file_path)
            root = tree.getroot()
            
            # Verify hashes
            if root.get('system_prompt_hash') != self.system_prompt_hash:
                print(f"System prompt hash mismatch in {file_path}")
            if root.get('action_schema_hash') != self.action_schema_hash:
                print(f"Action schema hash mismatch in {file_path}")
            
            # Get context and response
            context = root.find('context').text
            response = root.find('response').text
            
            if not context or not response:
                print(f"Missing context or response in {file_path}")
                return None
            
            # Create messages list
            messages = [
                {
                    "role": "system",
                    "content": self.system_prompt + self.action_schema
                },
                {
                    "role": "user",
                    "content": context
                },
                {
                    "role": "assistant", 
                    "content": response
                }
            ]
            
            return {"messages": messages}
            
        except Exception as e:
            print(f"Error processing {file_path}: {str(e)}")
            return None
    
    def create_dataset(self) -> int:
        """
        Create JSONL dataset from all valid iteration files
        
        Returns:
            Number of samples written to dataset
        """
        sample_count = 0
        
        # Create output directory if needed
        self.output_file.parent.mkdir(parents=True, exist_ok=True)
        
        with open(self.output_file, 'w', encoding='utf-8') as f:
            # Process each XML file in iterations directory
            for xml_file in sorted(self.iterations_dir.glob('*.xml')):
                sample = self._parse_iteration_file(xml_file)
                if sample:
                    json.dump(sample, f, ensure_ascii=False)
                    f.write('\n')
                    sample_count += 1
                    
        print(f"Created dataset with {sample_count} samples at {self.output_file}")
        return sample_count

def main():
    """Command line interface"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Create finetune dataset from iteration XML files')
    parser.add_argument('iterations_dir', type=str, help='Directory containing iteration XML files')
    parser.add_argument('system_prompt_file', type=str, help='Path to system prompt file')
    parser.add_argument('action_schema_file', type=str, help='Path to action schema file') 
    parser.add_argument('output_file', type=str, help='Path for output JSONL dataset')
    
    args = parser.parse_args()
    
    creator = FinetuneDatasetCreator(
        iterations_dir=args.iterations_dir,
        system_prompt_file=args.system_prompt_file,
        action_schema_file=args.action_schema_file,
        output_file=args.output_file
    )
    
    creator.create_dataset()

if __name__ == '__main__':
    main()