#!/usr/bin/env python3
"""
app.py – Quranic Data Training Pipeline Endpoint for ZeroGPU Spaces
--------------------------------------------------------------------
This script integrates a full Quranic data processing and training pipeline
into a Gradio interface endpoint. It is optimized for CPU-based training on 
Hugging Face ZeroGPU (using the Gradio SDK) and uses chunked incremental 
training, memory management, and gradient checkpointing to efficiently update 
Google's Gemma-2-2b model with Quranic data.

After every training chunk, the checkpoint folder is compressed into a ZIP file
and a download link is generated. Progress messages are streamed as the pipeline
runs.

Requirements:
 - Transformers (>=4.42.0)
 - Gradio (>=5.12.0)
 - PyTorch (==2.2.2)
 - psutil (==5.9.5)
 - Accelerate (>=0.26.0)
 - Hugging Face PRO subscription with ZeroGPU enabled (ensure your HF token is set as an environment variable HF_TOKEN)
 - Ubuntu CPU/Linux with access to ZeroGPU hardware via Spaces
 - Input data files placed in the project root.
 - Sufficient storage in "working_directory"

Author: [M-Saddam Hussain]
Date: March 2025
Data References: [Tanzil.net, IslamSource, QuranicCorpus]
"""

import json
import logging
import os
import traceback
import gc
import time
import psutil
import math
import zipfile
from datetime import datetime
from typing import Dict, List, Optional
from dataclasses import dataclass, asdict

import torch
# Limit PyTorch threads for CPU stability.
torch.set_num_threads(8)

from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    __version__ as transformers_version
)
from threading import Lock

import gradio as gr
import spaces

# Check for minimum required Transformers version for custom model support
MIN_TRANSFORMERS_VERSION = "4.42.0"
if tuple(map(int, transformers_version.split("."))) < tuple(map(int, MIN_TRANSFORMERS_VERSION.split("."))):
    logging.warning(f"Transformers version {transformers_version} detected. Please upgrade to at least {MIN_TRANSFORMERS_VERSION} for proper support of the 'gemma2' architecture.")

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('pipeline.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def manage_memory(threshold_percent: int = 90, min_available_mb: int = 500, sleep_duration: int = 10):
    """
    Check memory usage; if usage is high or available memory is low,
    force garbage collection and sleep briefly.
    """
    vm = psutil.virtual_memory()
    used_percent = vm.percent
    available_mb = vm.available / (1024 * 1024)
    logger.info(f"Memory usage: {used_percent}% used, {available_mb:.2f} MB available")
    if used_percent > threshold_percent or available_mb < min_available_mb:
        logger.warning("High memory usage detected, forcing garbage collection and sleeping...")
        gc.collect()
        time.sleep(sleep_duration)

@dataclass
class WordAnalysis:
    """Structured representation of word-level analysis"""
    arabic: str
    translation: str
    position: str
    morphology: Dict
    features: List[str]
    root: str
    location: str
    metadata: Dict

@dataclass
class VerseData:
    """Structured representation of verse-level data"""
    chapter: int
    verse: int
    arabic_text: str
    translation: str
    words: List[WordAnalysis]
    metadata: Dict

class QuranicDataset(Dataset):
    """Custom dataset for Quranic text training."""
    def __init__(self, processed_data: List[Dict], tokenizer):
        self.examples = []
        self.tokenizer = tokenizer
        for verse_data in processed_data:
            self.examples.extend(self._create_training_examples(verse_data))
    
    def _create_training_examples(self, verse_data: Dict) -> List[Dict]:
        examples = []
        text_block = (
            f"[VERSE {verse_data['chapter']}:{verse_data['verse']}]\n"
            f"Arabic: {verse_data['arabic_text']}\n"
            f"Translation: {verse_data['translation']}\n"
            "Morphological Analysis:\n"
        )
        for word in verse_data['words']:
            text_block += (
                f"[WORD] {word['arabic']}\n"
                f"Root: {word['root']}\n"
                f"Features: {', '.join(word['features'])}\n"
            )
        examples.append(self._format_example(text_block))
        return examples
    
    def _format_example(self, text: str) -> Dict:
        encodings = self.tokenizer(
            text,
            truncation=True,
            max_length=64,
            padding="max_length",
            return_tensors="pt"
        )
        return {
            "input_ids": encodings["input_ids"][0],
            "attention_mask": encodings["attention_mask"][0]
        }
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        return self.examples[idx]

class QuranicDataProcessor:
    """Processes Quranic data into structured training examples."""
    def __init__(self, source_dir: str, output_dir: str):
        self.source_dir = source_dir
        self.output_dir = output_dir
        self.morphological_data: Dict[str, Dict] = {}
        self.word_by_word_data: Dict[str, List[str]] = {}
        self.translation_data: Dict[str, str] = {}
        self.processing_lock = Lock()
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'json'), exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'txt'), exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'checkpoints'), exist_ok=True)
        logger.info(f"Initialized processor with source dir: {source_dir}")
    
    def load_source_files(self) -> bool:
        """Loads morphological, translation, and word-by-word data from project root."""
        try:
            logger.info("Loading morphological data...")
            morph_path = os.path.join(self.source_dir, 'quranic-corpus-morphology-0.4.txt')
            with open(morph_path, 'r', encoding='utf-8') as f:
                next(f)
                for line in f:
                    if line.strip() and not line.startswith('#'):
                        parts = line.strip().split('\t')
                        if len(parts) >= 4:
                            location = parts[0].strip('()')
                            self.morphological_data[location] = {
                                'form': parts[1],
                                'tag': parts[2],
                                'features': parts[3]
                            }
            logger.info(f"Loaded {len(self.morphological_data)} morphological entries")
            logger.info("Loading translation data...")
            trans_path = os.path.join(self.source_dir, 'en.sample.quran-maududi.txt')
            with open(trans_path, 'r', encoding='utf-8') as f:
                next(f)
                for line in f:
                    if line.strip():
                        parts = line.strip().split('|')
                        if len(parts) >= 3:
                            key = f"{parts[0]}:{parts[1]}"
                            self.translation_data[key] = parts[2].strip()
            logger.info(f"Loaded {len(self.translation_data)} verse translations")
            logger.info("Loading word-by-word data...")
            word_path = os.path.join(self.source_dir, 'en.w4w.qurandev.txt')
            with open(word_path, 'r', encoding='utf-8-sig') as f:
                lines = [line.strip() for line in f if line.strip()]
            sorted_keys = sorted(self.translation_data.keys(), key=lambda x: (int(x.split(':')[0]), int(x.split(':')[1])))
            if len(lines) != len(sorted_keys):
                logger.warning("Mismatch between word-by-word file and translation data")
            for i, verse_key in enumerate(sorted_keys):
                if i < len(lines):
                    words = [w.strip() for w in lines[i].split('|') if w.strip()]
                    self.word_by_word_data[verse_key] = words
            logger.info(f"Loaded word-by-word data for {len(self.word_by_word_data)} verses")
            return True
        except Exception as e:
            logger.error(f"Error loading source files: {str(e)}")
            logger.error(traceback.format_exc())
            return False

    def process_verse(self, chapter: int, verse: int) -> Optional[VerseData]:
        """Processes a single verse into structured format."""
        try:
            verse_ref = f"{chapter}:{verse}"
            logger.info(f"Processing verse {verse_ref}")
            translation = self.translation_data.get(verse_ref)
            if not translation:
                logger.warning(f"No translation for verse {verse_ref}")
                return None
            verse_word_list = self.word_by_word_data.get(verse_ref, [])
            if not verse_word_list:
                logger.warning(f"No word-by-word data for verse {verse_ref}")
                return None
            verse_words: List[WordAnalysis] = []
            arabic_text = ""
            for pos in range(1, len(verse_word_list) + 1):
                pattern = f"{chapter}:{verse}:{pos}:"
                matching_entries = [data for loc, data in self.morphological_data.items() if loc.startswith(pattern)]
                if not matching_entries:
                    logger.debug(f"No morphological data for {pattern}")
                    continue
                combined_form = " ".join(entry['form'] for entry in matching_entries)
                combined_features = []
                root = ""
                for entry in matching_entries:
                    features = entry['features'].split('|')
                    combined_features.extend(features)
                    if not root:
                        for f in features:
                            if 'ROOT:' in f:
                                root = f.split('ROOT:')[1]
                                break
                word_translation = verse_word_list[pos - 1]
                word = WordAnalysis(
                    arabic=combined_form,
                    translation=word_translation,
                    position=str(pos),
                    morphology=matching_entries[0],
                    features=combined_features,
                    root=root,
                    location=f"{chapter}:{verse}:{pos}",
                    metadata={}
                )
                verse_words.append(word)
                arabic_text += f" {combined_form}"
            verse_data = VerseData(
                chapter=chapter,
                verse=verse,
                arabic_text=arabic_text.strip(),
                translation=translation,
                words=verse_words,
                metadata={
                    "processed_timestamp": datetime.now().isoformat(),
                    "word_count": len(verse_words)
                }
            )
            self._save_verse_data(verse_data)
            return verse_data
        except Exception as e:
            logger.error(f"Error processing verse {chapter}:{verse}: {str(e)}")
            logger.error(traceback.format_exc())
            return None

    def _save_verse_data(self, verse_data: VerseData):
        """Saves processed verse data as JSON and TXT."""
        try:
            verse_ref = f"{verse_data.chapter}:{verse_data.verse}"
            json_path = os.path.join(self.output_dir, 'json', f'verse_{verse_ref.replace(":", "_")}.json')
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(asdict(verse_data), f, ensure_ascii=False, indent=2)
            txt_path = os.path.join(self.output_dir, 'txt', f'verse_{verse_ref.replace(":", "_")}.txt')
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(f"=== Verse {verse_ref} ===\n\n")
                f.write(f"Arabic Text:\n{verse_data.arabic_text}\n\n")
                f.write(f"Translation:\n{verse_data.translation}\n\n")
                f.write("Word Analysis:\n")
                for i, word in enumerate(verse_data.words, 1):
                    f.write(f"\nWord {i}:\n")
                    f.write(f"  Arabic: {word.arabic}\n")
                    f.write(f"  Translation: {word.translation}\n")
                    f.write(f"  Root: {word.root}\n")
                    f.write("  Features:\n")
                    for feature in word.features:
                        f.write(f"    - {feature}\n")
                    f.write("\n")
            logger.info(f"Saved verse data to {json_path} and {txt_path}")
        except Exception as e:
            logger.error(f"Error saving verse data: {str(e)}")
            logger.error(traceback.format_exc())

def zip_directory(folder_path: str, zip_path: str):
    """Compress the given folder into a zip file."""
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root_dir, _, files in os.walk(folder_path):
            for file in files:
                file_path = os.path.join(root_dir, file)
                arcname = os.path.relpath(file_path, folder_path)
                zipf.write(file_path, arcname)

class QuranicModelTrainer:
    """Trains the Gemma-2-2b model on Quranic data using chunked incremental updates."""
    def __init__(self,
                 model_name: str = "google/gemma-2-2b",
                 processed_data_dir: str = "processed_data",
                 checkpoint_dir: str = "checkpoints"):
        self.processed_data_dir = processed_data_dir
        self.checkpoint_dir = checkpoint_dir
        self.device = "cpu"  # Training on CPU; ZeroGPU will handle GPU access.
        logger.info("Loading tokenizer and model...")
        
        # Load tokenizer with additional special tokens and HF token from environment
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            token=os.environ.get("HF_TOKEN"),
            additional_special_tokens=["[VERSE]", "[WORD]", "[ROOT]", "[FEATURES]"],
            trust_remote_code=True
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
        
        # Load model using eager attention for Gemma2 and low_cpu_mem_usage.
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                token=os.environ.get("HF_TOKEN"),
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                attn_implementation="eager"
            )
        except Exception as e:
            logger.error(f"Error loading model directly: {str(e)}")
            logger.info("Attempting to load with fallback parameters...")
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained(
                model_name,
                token=os.environ.get("HF_TOKEN"),
                trust_remote_code=True
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                token=os.environ.get("HF_TOKEN"),
                config=config,
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                revision="main",
                attn_implementation="eager"
            )
        
        # Resize token embeddings to match tokenizer vocabulary size
        self.model.resize_token_embeddings(len(self.tokenizer))
        self.model.train()
        self.model.config.use_cache = False
        
        if hasattr(self.model, "gradient_checkpointing_enable"):
            self.model.gradient_checkpointing_enable()
        else:
            logger.warning("Gradient checkpointing not available for this model")
    
    def prepare_training_data(self, chapter_data: List[Dict]) -> Dataset:
        """Creates a QuranicDataset from processed chapter data."""
        return QuranicDataset(chapter_data, self.tokenizer)
    
    def train_chapter(self,
                      chapter_num: int,
                      processed_verses: List[Dict],
                      chunk_size: int = 5,       # Reduced chunk size to help with memory
                      num_train_epochs: int = 5,   # Fewer epochs for testing
                      per_device_train_batch_size: int = 1,
                      learning_rate: float = 3e-5,
                      weight_decay: float = 0.01,
                      gradient_accumulation_steps: int = 32):
        """
        Splits chapter data into chunks and trains incrementally.
        Yields progress messages and download links for each checkpoint.
        """
        total_examples = len(processed_verses)
        total_chunks = math.ceil(total_examples / chunk_size)
        yield f"Chapter {chapter_num}: {total_examples} examples, {total_chunks} chunks."
        for chunk_index in range(total_chunks):
            chunk_data = processed_verses[chunk_index * chunk_size: (chunk_index + 1) * chunk_size]
            dataset = self.prepare_training_data(chunk_data)
            chunk_output_dir = os.path.join(self.checkpoint_dir, f"chapter_{chapter_num}", f"chunk_{chunk_index}")
            os.makedirs(chunk_output_dir, exist_ok=True)
            training_args = TrainingArguments(
                output_dir=chunk_output_dir,
                overwrite_output_dir=True,
                num_train_epochs=num_train_epochs,
                per_device_train_batch_size=per_device_train_batch_size,
                learning_rate=learning_rate,
                weight_decay=weight_decay,
                gradient_accumulation_steps=gradient_accumulation_steps,
                fp16=False,
                remove_unused_columns=False,
                logging_steps=50,
                report_to="none",
                evaluation_strategy="no",
                use_cpu=True,  # Using CPU flag (replaces no_cuda)
                dataloader_num_workers=0,
                dataloader_pin_memory=False
            )
            data_collator = DataCollatorForLanguageModeling(
                tokenizer=self.tokenizer,
                mlm=False
            )
            trainer = Trainer(
                model=self.model,
                args=training_args,
                train_dataset=dataset,
                tokenizer=self.tokenizer,
                data_collator=data_collator
            )
            yield f"Training chunk {chunk_index+1}/{total_chunks} for Chapter {chapter_num}..."
            trainer.train()
            trainer.save_model(chunk_output_dir)
            # Clean up trainer and dataset
            del trainer, dataset
            gc.collect()
            manage_memory()
            # Compress the checkpoint folder into a zip file
            zip_path = f"{chunk_output_dir}.zip"
            zip_directory(chunk_output_dir, zip_path)
            # Generate a download link (assuming working_directory is served)
            download_link = f"<a href='{zip_path}' target='_blank'>Download Chapter {chapter_num} Chunk {chunk_index+1} Checkpoint</a>"
            yield f"Checkpoint saved for Chapter {chapter_num} Chunk {chunk_index+1}: {download_link}"
        yield f"Completed training for Chapter {chapter_num}."
    
class QuranicPipeline:
    """Integrates data processing and incremental model training for all chapters."""
    def __init__(self,
                 source_dir: str = ".",
                 working_dir: str = "working_directory",
                 start_chapter: int = 1,
                 end_chapter: int = 114):
        self.source_dir = source_dir
        self.working_dir = working_dir
        self.start_chapter = start_chapter
        self.end_chapter = end_chapter
        self.setup_directories()
        global logger
        logger = logging.getLogger(__name__)
        self.state = {
            "last_processed_chapter": 0,
            "last_trained_chapter": 0,
            "current_state": "initialized",
            "errors": [],
            "start_time": datetime.now().isoformat()
        }
        self.load_state()
        try:
            logger.info("Initializing Quranic Data Processor...")
            self.processor = QuranicDataProcessor(
                source_dir=self.source_dir,
                output_dir=os.path.join(self.working_dir, "processed_data")
            )
            logger.info("Initializing Quranic Model Trainer...")
            self.trainer = QuranicModelTrainer(
                model_name="google/gemma-2-2b",
                processed_data_dir=os.path.join(self.working_dir, "processed_data"),
                checkpoint_dir=os.path.join(self.working_dir, "checkpoints")
            )
            self.state["current_state"] = "ready"
            self.save_state()
        except Exception as e:
            self.handle_error("Initialization failed", e)
            raise

    def setup_directories(self):
        dirs = [
            self.working_dir,
            os.path.join(self.working_dir, "processed_data"),
            os.path.join(self.working_dir, "checkpoints"),
            os.path.join(self.working_dir, "logs"),
            os.path.join(self.working_dir, "state")
        ]
        for d in dirs:
            os.makedirs(d, exist_ok=True)

    def load_state(self):
        state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
        if os.path.exists(state_file):
            try:
                with open(state_file, 'r') as f:
                    saved_state = json.load(f)
                self.state.update(saved_state)
                logger.info(f"Loaded previous state: Last processed chapter {self.state.get('last_processed_chapter')}, "
                            f"last trained chapter {self.state.get('last_trained_chapter')}")
            except Exception as e:
                logger.warning(f"Could not load previous state: {str(e)}")

    def save_state(self):
        state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
        with open(state_file, 'w') as f:
            json.dump(self.state, f, indent=2)

    def handle_error(self, context: str, error: Exception):
        error_detail = {
            "timestamp": datetime.now().isoformat(),
            "context": context,
            "error": str(error),
            "traceback": traceback.format_exc()
        }
        self.state.setdefault("errors", []).append(error_detail)
        logger.error(f"{context}: {str(error)}")
        self.save_state()

    def run_pipeline(self):
        """
        Runs processing and training for chapters sequentially.
        Yields progress messages and download links as checkpoints become available.
        Finally, saves the final model and yields its download link.
        """
        yield "Starting pipeline execution..."
        if not self.processor.load_source_files():
            yield "Failed to load source files. Stopping pipeline."
            return
        for chapter in range(self.start_chapter, self.end_chapter + 1):
            yield f"=== Processing Chapter {chapter} ==="
            processed_chapter_data = []
            verse = 1
            while True:
                verse_data = self.processor.process_verse(chapter, verse)
                if verse_data is None:
                    break
                processed_chapter_data.append(asdict(verse_data))
                verse += 1
            if processed_chapter_data:
                # Iterate over checkpoint messages from training this chapter
                for msg in self.trainer.train_chapter(chapter, processed_chapter_data):
                    yield msg
                self.state["last_trained_chapter"] = chapter
                self.save_state()
            else:
                yield f"No processed data for Chapter {chapter}."
            self.state["last_processed_chapter"] = chapter
            self.save_state()
            manage_memory()
        yield "Pipeline execution completed."
        # Save the final model and tokenizer
        final_model_dir = os.path.join(self.working_dir, "final_model")
        os.makedirs(final_model_dir, exist_ok=True)
        self.trainer.model.save_pretrained(final_model_dir)
        self.trainer.tokenizer.save_pretrained(final_model_dir)
        # Compress the final model folder into a zip
        final_zip = f"{final_model_dir}.zip"
        zip_directory(final_model_dir, final_zip)
        final_link = f"<a href='{final_zip}' target='_blank'>Download Final Model</a>"
        yield f"Final model saved. {final_link}"

@spaces.GPU()  # Request ZeroGPU hardware for the Space
def start_pipeline():
    try:
        yield "Starting Quranic Training Pipeline with Gemma-2-2b"
        yield f"PyTorch version: {torch.__version__}"
        yield f"CUDA available: {torch.cuda.is_available()}"
        if torch.cuda.is_available():
            yield f"CUDA device count: {torch.cuda.device_count()}"
            yield f"CUDA device name: {torch.cuda.get_device_name(0)}"
        if not os.environ.get("HF_TOKEN"):
            yield "WARNING: HF_TOKEN environment variable not set. Model loading may fail."
        required_files = [
            'quranic-corpus-morphology-0.4.txt',
            'en.sample.quran-maududi.txt',
            'en.w4w.qurandev.txt'
        ]
        missing_files = [f for f in required_files if not os.path.exists(f)]
        if missing_files:
            yield f"Missing required data files: {', '.join(missing_files)}"
            return
        pipeline = QuranicPipeline(
            source_dir=".",
            working_dir="working_directory",
            start_chapter=1,
            end_chapter=114
        )
        for msg in pipeline.run_pipeline():
            yield msg
    except Exception as e:
        error_msg = f"Pipeline execution failed: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        yield error_msg

iface = gr.Interface(
    fn=start_pipeline,
    inputs=[],
    outputs=gr.HTML(label="Pipeline Status"),
    title="Quranic Training Pipeline for Gemma-2-2b",
    description="""This pipeline fine-tunes Google's Gemma-2-2b model on Quranic data.
    
Click 'Submit' to trigger the Quranic data processing and training pipeline on ZeroGPU.
    
During execution, download links for each checkpoint (and the final model) will be provided as they become available.
    
Requirements:
- Transformers (>=4.42.0)
- Gradio (>=5.12.0)
- PyTorch (==2.2.2)
- psutil (==5.9.5)
- Accelerate (>=0.26.0)
    
The pipeline processes all 114 chapters of the Quran sequentially, with memory management optimizations for ZeroGPU environments."""
)

if __name__ == "__main__":
    iface.launch()