"""
TTS Dataset Collection Tool with Font Support and Enhanced Error Handling
"""

import os
import json
import nltk
import gradio as gr
from datetime import datetime
from pathlib import Path
import shutil
import logging
from typing import Dict, List, Tuple, Optional
import traceback

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Font configurations
FONT_STYLES = {
    "english_serif": {
        "name": "Times New Roman",
        "family": "serif",
        "css": "font-family: 'Times New Roman', serif;"
    },
    "english_sans": {
        "name": "Arial",
        "family": "sans-serif",
        "css": "font-family: Arial, sans-serif;"
    },
    "nastaliq": {
        "name": "Nastaliq",
        "family": "Jameel Noori Nastaleeq",
        "css": "font-family: 'Jameel Noori Nastaleeq', serif;"
    },
    "naskh": {
        "name": "Naskh",
        "family": "Traditional Arabic",
        "css": "font-family: 'Traditional Arabic', serif;"
    }
}

class TTSDatasetCollector:
    """Manages TTS dataset collection and organization with enhanced features"""
    
    def __init__(self):
        """Initialize the TTS Dataset Collector"""
        # Initialize NLTK
        self._initialize_nltk()
        
        # Set up paths and directories
        self.root_path = Path(os.path.dirname(os.path.abspath(__file__))) / "dataset"
        self.sentences: List[str] = []
        self.current_index: int = 0
        self.current_font: str = "english_serif"
        self.setup_directories()
        
        logger.info("TTS Dataset Collector initialized")
    
    def _initialize_nltk(self) -> None:
        """Initialize NLTK with error handling"""
        try:
            nltk.download('punkt', quiet=True)
            logger.info("NLTK punkt tokenizer downloaded successfully")
        except Exception as e:
            logger.error(f"Failed to download NLTK data: {str(e)}")
            logger.error(traceback.format_exc())
            raise RuntimeError("Failed to initialize NLTK. Please check your internet connection.")
    
    def setup_directories(self) -> None:
        """Create necessary directory structure with logging"""
        try:
            # Create main dataset directory
            self.root_path.mkdir(exist_ok=True)
            
            # Create subdirectories
            for subdir in ['audio', 'transcriptions', 'metadata', 'fonts']:
                (self.root_path / subdir).mkdir(exist_ok=True)
            
            # Initialize log file
            log_file = self.root_path / 'dataset_log.txt'
            if not log_file.exists():
                with open(log_file, 'w', encoding='utf-8') as f:
                    f.write(f"Dataset collection initialized on {datetime.now().isoformat()}\n")
            
            logger.info("Directory structure created successfully")
            
        except Exception as e:
            logger.error(f"Failed to create directory structure: {str(e)}")
            logger.error(traceback.format_exc())
            raise RuntimeError("Failed to initialize directory structure")
    
    def log_operation(self, message: str, level: str = "info") -> None:
        """Log operations with timestamp and level"""
        try:
            log_file = self.root_path / 'dataset_log.txt'
            timestamp = datetime.now().isoformat()
            
            with open(log_file, 'a', encoding='utf-8') as f:
                f.write(f"[{timestamp}] [{level.upper()}] {message}\n")
            
            if level.lower() == "error":
                logger.error(message)
            else:
                logger.info(message)
                
        except Exception as e:
            logger.error(f"Failed to log operation: {str(e)}")
    
    def load_text_file(self, file) -> Tuple[bool, str]:
        """Process and load text file with enhanced error handling"""
        if not file:
            return False, "No file provided"
        
        try:
            # Validate file extension
            if not file.name.endswith('.txt'):
                return False, "Only .txt files are supported"
            
            with open(file.name, 'r', encoding='utf-8') as f:
                text = f.read()
            
            # Validate text content
            if not text.strip():
                return False, "File is empty"
            
            # Tokenize sentences
            self.sentences = nltk.sent_tokenize(text)
            if not self.sentences:
                return False, "No valid sentences found in file"
            
            self.current_index = 0
            
            # Log success
            self.log_operation(
                f"Loaded text file: {file.name} with {len(self.sentences)} sentences"
            )
            
            return True, f"Successfully loaded {len(self.sentences)} sentences"
            
        except UnicodeDecodeError:
            error_msg = "File encoding error. Please ensure the file is UTF-8 encoded"
            self.log_operation(error_msg, "error")
            return False, error_msg
        except Exception as e:
            error_msg = f"Error loading file: {str(e)}"
            self.log_operation(error_msg, "error")
            logger.error(traceback.format_exc())
            return False, error_msg
    
    # Remaining methods go here ...

def create_interface():
    """Create Gradio interface with enhanced features"""
    
    # Create custom CSS for fonts
    custom_css = """
    .gradio-container {
        max-width: 1200px !important;
    }
    .record-button {
        font-size: 1.2em !important;
        padding: 20px !important;
    }
    """
    
    # Add font-face declarations
    for font_style, font_info in FONT_STYLES.items():
        if font_style in ['nastaliq', 'naskh']:
            custom_css += f"""
            @font-face {{
                font-family: '{font_info["family"]}';
                src: url('fonts/{font_info["family"]}.ttf') format('truetype');
            }}
            """
    
    collector = TTSDatasetCollector()
    
    with gr.Blocks(title="TTS Dataset Collection Tool", css=custom_css) as interface:
        gr.Markdown("# TTS Dataset Collection Tool")
        
        with gr.Row():
            # Left column - Configuration
            with gr.Column():
                file_input = gr.File(
                    label="Upload Text File (.txt)",
                    file_types=[".txt"]
                )
                speaker_id = gr.Textbox(
                    label="Speaker ID",
                    placeholder="Enter unique speaker identifier (letters and numbers only)"
                )
                dataset_name = gr.Textbox(
                    label="Dataset Name",
                    placeholder="Enter dataset name (letters and numbers only)"
                )
                font_select = gr.Dropdown(
                    choices=list(FONT_STYLES.keys()),
                    value="english_serif",
                    label="Select Font Style"
                )
            
            # Right column - Recording
            with gr.Column():
                current_text = gr.HTML(
                    label="Current Sentence"
                )
                audio_recorder = gr.Audio(
                    label="Record Audio",
                    type="filepath"
                )
                next_text = gr.HTML(
                    label="Next Sentence"
                )
        
        # Controls
        with gr.Row():
            prev_btn = gr.Button("Previous", variant="secondary")
            next_btn = gr.Button("Next", variant="secondary")
            save_btn = gr.Button("Save Recording", variant="primary")
        
        # Status and Progress
        with gr.Row():
            progress = gr.Textbox(
                label="Progress",
                interactive=False
            )
            status = gr.Textbox(
                label="Status",
                interactive=False,
                max_lines=3
            )
        
        # Dataset Info
        with gr.Row():
            dataset_info = gr.JSON(
                label="Dataset Statistics",
                value={}
            )
        
        def update_font(font_style):
            """Update font and refresh display"""
            success, msg = collector.set_font(font_style)
            if not success:
                return {status: msg}
            
            return update_display()
        
        def load_file(file):
            """Handle file loading with enhanced error reporting"""
            if not file:
                return {
                    current_text: "",
                    next_text: "",
                    progress: "",
                    status: "⚠️ No file selected"
                }

            success, msg = collector.load_text_file(file)
            if not success:
                return {
                    current_text: "",
                    next_text: "",
                    progress: "",
                    status: f"❌ {msg}",
                    dataset_info: update_dataset_info()
                }
            
            return {
                current_text: collector.get_styled_text(collector.sentences[0]),
                next_text: collector.get_styled_text(collector.sentences[1]) if len(collector.sentences) > 1 else "",
                progress: f"📖 Sentence 1 of {len(collector.sentences)}",
                status: f"✅ {msg}",
                dataset_info: update_dataset_info()
            }
        
        # Remaining methods and event handlers go here ...

        return interface

if __name__ == "__main__":
    try:
        interface = create_interface()
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            enable_queue=True
        )
    except Exception as e:
        logger.error(f"Failed to launch interface: {str(e)}")
        logger.error(traceback.format_exc())
        raise