"""
End-to-End Afaan Oromoo NLP Corpus Pipeline
==========================================

Builds, normalizes, tokenizes, stems, and saves corpus data.
"""

import json
from pathlib import Path
from typing import Dict

from qubee_nlp.corpus.afaan_oromoo_nlp_corpus import AfaanOromooCorpusBuilder
from qubee_nlp.tokenizer import Tokenizer
from qubee_nlp.stemmer import Stemmer
from qubee_nlp.normalizer import Normalizer


class AfaanOromooCorpusPipeline:
    """
    High-level pipeline that integrates:
    Corpus Builder → Normalizer → Tokenizer → Stemmer
    """

    def __init__(self, output_dir: str = "corpus_output"):
        self.base_dir = Path(output_dir)
        self.raw_dir = self.base_dir / "raw"
        self.processed_dir = self.base_dir / "processed"

        self.raw_dir.mkdir(parents=True, exist_ok=True)
        self.processed_dir.mkdir(parents=True, exist_ok=True)

        self.builder = AfaanOromooCorpusBuilder(output_dir=self.raw_dir)
        self.normalizer = Normalizer()
        self.tokenizer = Tokenizer()
        self.stemmer = Stemmer()

    # ===============================
    # CORE PIPELINE
    # ===============================

    def build_raw_corpus(self) -> Dict:
        corpus = self.builder.run_complete_pipeline()

        path = self.raw_dir / "raw_corpus.json"
        self._save_json(corpus, path)

        return corpus

    def process_corpus(self, raw_corpus: Dict) -> Dict:
        processed = {}

        for category, articles in raw_corpus.items():
            processed[category] = []

            for article in articles:
                text = article.get("content", "")

                normalized = self.normalizer.normalize(text)
                tokens = self.tokenizer.tokenize(normalized)
                stems = [self.stemmer.stem(t) for t in tokens]

                processed[category].append({
                    "title": article.get("title"),
                    "source": article.get("source"),
                    "original_text": text,
                    "normalized_text": normalized,
                    "tokens": tokens,
                    "stems": stems,
                })

        path = self.processed_dir / "processed_corpus.json"
        self._save_json(processed, path)

        return processed

    def run(self) -> Dict:
        raw = self.build_raw_corpus()
        return self.process_corpus(raw)

    # ===============================
    # UTIL
    # ===============================

    @staticmethod
    def _save_json(data: Dict, path: Path):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
