"""
Data management and loading utilities for Afaan Oromoo NLP.
"""

import os
import json
import pickle
from typing import Dict, List, Any, Optional, Union
import csv

# Data directory
_DATA_DIR = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'data')


def ensure_data_dir() -> str:
    """
    Ensure data directory exists.
    
    Returns:
        Path to data directory
    """
    os.makedirs(_DATA_DIR, exist_ok=True)
    return _DATA_DIR


def save_json(data: Any, filename: str) -> str:
    """
    Save data as JSON file.
    
    Args:
        data: Data to save
        filename: Name of file (without .json extension)
        
    Returns:
        Path to saved file
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.json")
    
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    
    return filepath


def load_json(filename: str) -> Any:
    """
    Load data from JSON file.
    
    Args:
        filename: Name of file (without .json extension)
        
    Returns:
        Loaded data
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.json")
    
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"Data file not found: {filepath}")
    
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)


def save_pickle(data: Any, filename: str) -> str:
    """
    Save data as pickle file.
    
    Args:
        data: Data to save
        filename: Name of file (without .pkl extension)
        
    Returns:
        Path to saved file
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.pkl")
    
    with open(filepath, 'wb') as f:
        pickle.dump(data, f)
    
    return filepath


def load_pickle(filename: str) -> Any:
    """
    Load data from pickle file.
    
    Args:
        filename: Name of file (without .pkl extension)
        
    Returns:
        Loaded data
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.pkl")
    
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"Pickle file not found: {filepath}")
    
    with open(filepath, 'rb') as f:
        return pickle.load(f)


def save_csv(data: List[List[Any]], filename: str, headers: Optional[List[str]] = None) -> str:
    """
    Save data as CSV file.
    
    Args:
        data: List of rows (each row is a list)
        filename: Name of file (without .csv extension)
        headers: Optional column headers
        
    Returns:
        Path to saved file
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.csv")
    
    with open(filepath, 'w', encoding='utf-8', newline='') as f:
        writer = csv.writer(f)
        if headers:
            writer.writerow(headers)
        writer.writerows(data)
    
    return filepath


def load_csv(filename: str, has_headers: bool = True) -> Union[List[List[str]], Dict[str, List[str]]]:
    """
    Load data from CSV file.
    
    Args:
        filename: Name of file (without .csv extension)
        has_headers: Whether the CSV has headers
        
    Returns:
        If has_headers: Dictionary with column names as keys
        If no headers: List of rows
    """
    data_dir = ensure_data_dir()
    filepath = os.path.join(data_dir, f"{filename}.csv")
    
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"CSV file not found: {filepath}")
    
    with open(filepath, 'r', encoding='utf-8') as f:
        reader = csv.reader(f)
        
        if has_headers:
            headers = next(reader)
            data = {header: [] for header in headers}
            
            for row in reader:
                for i, value in enumerate(row):
                    if i < len(headers):
                        data[headers[i]].append(value)
            
            return data
        else:
            return [row for row in reader]


class DatasetManager:
    """
    Manager for Afaan Oromoo NLP datasets.
    """
    
    def __init__(self, dataset_name: str):
        """
        Initialize dataset manager.
        
        Args:
            dataset_name: Name of dataset
        """
        self.dataset_name = dataset_name
        self.data_dir = ensure_data_dir()
        self.dataset_dir = os.path.join(self.data_dir, dataset_name)
        os.makedirs(self.dataset_dir, exist_ok=True)
    
    def save_dataset(self, 
                     train_data: Any,
                     val_data: Optional[Any] = None,
                     test_data: Optional[Any] = None,
                     metadata: Optional[Dict[str, Any]] = None) -> None:
        """
        Save dataset splits.
        
        Args:
            train_data: Training data
            val_data: Validation data (optional)
            test_data: Test data (optional)
            metadata: Dataset metadata (optional)
        """
        # Save splits
        self._save_split(train_data, 'train')
        if val_data is not None:
            self._save_split(val_data, 'val')
        if test_data is not None:
            self._save_split(test_data, 'test')
        
        # Save metadata
        if metadata is not None:
            metadata_file = os.path.join(self.dataset_dir, 'metadata.json')
            with open(metadata_file, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)
    
    def load_dataset(self) -> Dict[str, Any]:
        """
        Load dataset.
        
        Returns:
            Dictionary with dataset splits and metadata
        """
        dataset = {}
        
        # Load splits
        for split in ['train', 'val', 'test']:
            split_file = os.path.join(self.dataset_dir, f'{split}.pkl')
            if os.path.exists(split_file):
                with open(split_file, 'rb') as f:
                    dataset[split] = pickle.load(f)
        
        # Load metadata
        metadata_file = os.path.join(self.dataset_dir, 'metadata.json')
        if os.path.exists(metadata_file):
            with open(metadata_file, 'r', encoding='utf-8') as f:
                dataset['metadata'] = json.load(f)
        
        return dataset
    
    def _save_split(self, data: Any, split_name: str) -> None:
        """Save a dataset split."""
        split_file = os.path.join(self.dataset_dir, f'{split_name}.pkl')
        with open(split_file, 'wb') as f:
            pickle.dump(data, f)
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get dataset statistics.
        
        Returns:
            Dictionary with dataset statistics
        """
        dataset = self.load_dataset()
        stats = {}
        
        for split_name, split_data in dataset.items():
            if split_name == 'metadata':
                continue
            
            if isinstance(split_data, list):
                stats[split_name] = {
                    'samples': len(split_data),
                    'type': type(split_data[0]).__name__ if split_data else 'None'
                }
        
        return stats


class CorpusLoader:
    """
    Loader for Afaan Oromoo text corpora.
    """
    
    def __init__(self):
        """Initialize corpus loader."""
        self.data_dir = ensure_data_dir()
    
    def load_text_corpus(self, corpus_name: str) -> List[str]:
        """
        Load text corpus from file.
        
        Args:
            corpus_name: Name of corpus file (without extension)
            
        Returns:
            List of text documents
        """
        filepath = os.path.join(self.data_dir, f'{corpus_name}.txt')
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"Corpus file not found: {filepath}")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            # Split by empty lines or other delimiters
            content = f.read()
            documents = [doc.strip() for doc in content.split('\n\n') if doc.strip()]
        
        return documents
    
    def save_text_corpus(self, documents: List[str], corpus_name: str) -> str:
        """
        Save text corpus to file.
        
        Args:
            documents: List of text documents
            corpus_name: Name for corpus file
            
        Returns:
            Path to saved file
        """
        filepath = os.path.join(self.data_dir, f'{corpus_name}.txt')
        
        with open(filepath, 'w', encoding='utf-8') as f:
            for doc in documents:
                f.write(doc.strip())
                f.write('\n\n')
        
        return filepath
    
    def load_tagged_corpus(self, corpus_name: str) -> List[List[tuple]]:
        """
        Load POS-tagged corpus.
        
        Args:
            corpus_name: Name of tagged corpus
            
        Returns:
            List of sentences, each sentence is list of (word, tag) tuples
        """
        filepath = os.path.join(self.data_dir, f'{corpus_name}_tagged.json')
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"Tagged corpus not found: {filepath}")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Convert to list of sentences with tuples
        tagged_corpus = []
        for sentence in data:
            tagged_sentence = [(item['word'], item['tag']) for item in sentence]
            tagged_corpus.append(tagged_sentence)
        
        return tagged_corpus
    
    def save_tagged_corpus(self, 
                          tagged_corpus: List[List[tuple]], 
                          corpus_name: str) -> str:
        """
        Save POS-tagged corpus.
        
        Args:
            tagged_corpus: Tagged corpus data
            corpus_name: Name for corpus
            
        Returns:
            Path to saved file
        """
        filepath = os.path.join(self.data_dir, f'{corpus_name}_tagged.json')
        
        # Convert to JSON-serializable format
        data = []
        for sentence in tagged_corpus:
            sentence_data = [
                {'word': word, 'tag': tag}
                for word, tag in sentence
            ]
            data.append(sentence_data)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        return filepath


# Create data directory if it doesn't exist
ensure_data_dir()


__all__ = [
    'save_json',
    'load_json',
    'save_pickle',
    'load_pickle',
    'save_csv',
    'load_csv',
    'DatasetManager',
    'CorpusLoader',
    'ensure_data_dir',
]