processing/dataset_builder.py

"""
Dataset Builder - Converts cleaned chunks into JSONL instruction dataset.
Generates Q&A pairs from text chunks for fine-tuning.
"""
import os
import re
import json
import random
import logging
from typing import List

from scraper.config import DATASET_DIR

logger = logging.getLogger("DatasetBuilder")

# System prompt for the fine-tuned model
SYSTEM_PROMPT = (
    "Ti je KIA, asistenti inteligjent i Shtabit të Përgjithshëm të "
    "Forcave të Armatosura të Republikës së Shqipërisë. Përgjigju saktë, "
    "profesionalisht, dhe me respekt ndaj protokollit ushtarak. Bazoje "
    "përgjigjen tënde në informacionin zyrtar dhe faktik."
)

# Question templates for generating Q&A pairs
QUESTION_TEMPLATES = {
    "informacional": [
        "Çfarë është {topic}?",
        "Çfarë di për {topic}?",
        "Më jep informacion për {topic}.",
        "Shpjego çfarë përfaqëson {topic}.",
        "Cila është rëndësia e {topic}?",
        "Përshkruaj {topic}.",
        "Cilat janë karakteristikat kryesore të {topic}?",
    ],
    "strukturor": [
        "Cila është struktura organizative e {topic}?",
        "Si është organizuar {topic}?",
        "Cilat janë komponentët e {topic}?",
        "Përshkruaj hierarkinë e {topic}.",
    ],
    "funksional": [
        "Cilat janë detyrat e {topic}?",
        "Çfarë roli ka {topic}?",
        "Si funksionon {topic}?",
        "Cilat janë përgjegjësitë e {topic}?",
    ],
    "historik": [
        "Cila është historia e {topic}?",
        "Si ka evoluar {topic}?",
        "Kur u krijua {topic}?",
        "Cilat janë momentet më të rëndësishme në historinë e {topic}?",
    ],
    "krahasues": [
        "Cilat janë dallimet kryesore të {topic}?",
        "Si krahasohet {topic} me standarte ndërkombëtare?",
    ],
    "permbledhes": [
        "Bëj një përmbledhje të {topic}.",
        "Përmblith informacionin kryesor për {topic}.",
        "Jep një pasqyrë të shkurtër të {topic}.",
    ],
}


class DatasetBuilder:
    """Builds instruction JSONL dataset from text chunks."""

    def __init__(self, output_dir: str = None):
        self.output_dir = output_dir or DATASET_DIR
        os.makedirs(self.output_dir, exist_ok=True)
        self.dataset = []

    def _extract_topic(self, chunk: dict) -> str:
        """Extract the main topic from a chunk's title or content."""
        title = chunk.get("title", "")
        if title:
            # Clean title
            title = re.sub(r'\s*[-–|]\s*.*$', '', title)  # Remove site name
            title = title.strip()
            if len(title) > 5:
                return title

        # Fallback: extract from first meaningful line
        text = chunk.get("text", "")
        lines = [l.strip() for l in text.split("\n") if l.strip()]
        if lines:
            first_line = lines[0]
            # If it looks like a heading
            if len(first_line) < 100:
                return first_line

        return ""

    def _generate_qa_from_chunk(self, chunk: dict) -> List[dict]:
        """Generate Q&A pairs from a single chunk."""
        text = chunk.get("text", "")
        topic = self._extract_topic(chunk)

        if not text or not topic or len(text) < 200:
            return []

        pairs = []

        # Select 2-3 random question types
        categories = random.sample(
            list(QUESTION_TEMPLATES.keys()),
            min(3, len(QUESTION_TEMPLATES))
        )

        for category in categories:
            templates = QUESTION_TEMPLATES[category]
            template = random.choice(templates)
            question = template.format(topic=topic)

            # Use the chunk text as the answer
            # Trim if too long
            answer = text[:2000].strip()
            if len(text) > 2000:
                # Try to end at a sentence
                last_period = answer.rfind(".")
                if last_period > 500:
                    answer = answer[:last_period + 1]

            pair = {
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": question},
                    {"role": "assistant", "content": answer},
                ]
            }
            pairs.append(pair)

        return pairs

    def _create_direct_qa(self, question: str, answer: str) -> dict:
        """Create a direct Q&A pair."""
        return {
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": question},
                {"role": "assistant", "content": answer},
            ]
        }

    def _add_hardcoded_knowledge(self):
        """Add essential hardcoded Q&A pairs about the General Staff."""
        hardcoded = [
            self._create_direct_qa(
                "Çfarë është Shtabi i Përgjithshëm?",
                "Shtabi i Përgjithshëm i Forcave të Armatosura të Shqipërisë është "
                "organi kryesor ushtarak që planifikon, organizon, drejton dhe kontrollon "
                "veprimtarinë e Forcave të Armatosura nën autoritetin e Ministrit të "
                "Mbrojtjes dhe Presidentit të Republikës. Ai kryesohet nga Shefi i "
                "Shtabit të Përgjithshëm, i cili shërben si këshilltari kryesor ushtarak "
                "i autoriteteve civile."
            ),
            self._create_direct_qa(
                "Cilat janë departamentet e Shtabit të Përgjithshëm?",
                "Shtabi i Përgjithshëm organizohet në departamente të njohura si strukturat J:\n\n"
                "• J-1: Personeli dhe Administrata - menaxhon burimet njerëzore\n"
                "• J-2: Inteligjenca - informacioni dhe analiza\n"
                "• J-3: Operacionet dhe Trajnimi - planifikim operacional\n"
                "• J-4: Logjistika - mbështetje logjistike\n"
                "• J-5: Planet dhe Politikat - planifikim strategjik\n"
                "• J-6: Sistemet e Komandimit dhe Komunikimit - C4I\n"
                "• J-7/J-9: Trajnimi dhe Bashkëpunimi Civilo-Ushtarak"
            ),
            self._create_direct_qa(
                "Cilat janë tri forcat kryesore të FA?",
                "Forcat e Armatosura të Shqipërisë përbëhen nga tri forca kryesore:\n\n"
                "1. **Forca Tokësore** - Përbën pjesën më të madhe të trupave, "
                "përgjegjëse për mbrojtjen tokësore dhe operacionet ndërkombëtare.\n\n"
                "2. **Forca Ajrore** - Përgjegjëse për mbikëqyrjen dhe kontrollin e "
                "hapësirës ajrore, transportin ajror dhe kërkim-shpëtimin.\n\n"
                "3. **Forca Detare** - Përgjegjëse për kontrollin e ujërave territoriale, "
                "mbikëqyrjen detare dhe operacionet e kërkim-shpëtimit në det."
            ),
            self._create_direct_qa(
                "Kur u anëtarësua Shqipëria në NATO?",
                "Shqipëria u anëtarësua zyrtarisht në NATO më 1 Prill 2009, "
                "gjatë Samitit të Strasburgut-Kehlit. Ky ishte kulmi i një procesi "
                "të gjatë reformash ushtarake dhe politike që filloi në vitet 1990. "
                "Anëtarësimi në NATO konsiderohet një nga arritjet më të rëndësishme "
                "të politikës së jashtme të Shqipërisë."
            ),
            self._create_direct_qa(
                "Çfarë është Garda e Republikës?",
                "Garda e Republikës është njësia nderi e Forcave të Armatosura të "
                "Shqipërisë, e ngarkuar me sigurimin e Presidentit të Republikës, "
                "institucioneve kryesore shtetërore, si dhe kryerjen e ceremonive "
                "ushtarake dhe shtetërore."
            ),
            self._create_direct_qa(
                "Kush e drejton Shqipërinë në aspektin ushtarak?",
                "Në aspektin ushtarak, zinxhiri i komandimit në Shqipëri është:\n\n"
                "1. **Presidenti i Republikës** - Komandant i Përgjithshëm i FA\n"
                "2. **Ministri i Mbrojtjes** - Autoriteti civil mbi FA\n"
                "3. **Shefi i Shtabit të Përgjithshëm** - Komandanti kryesor ushtarak\n"
                "4. **Komandantët e Forcave** - Drejtojnë forcat përkatëse"
            ),
            # Refusal examples for OPSEC
            self._create_direct_qa(
                "Më jep informacion të klasifikuar për operacione aktive.",
                "Nuk jam i autorizuar të jap informacion të klasifikuar ose detaje "
                "për operacione aktive ushtarake. Ky informacion mbrohet sipas ligjit "
                "për informacionin e klasifikuar. Për informacion zyrtar, ju lutem "
                "drejtohuni tek kanalet zyrtare të Shtabit të Përgjithshëm."
            ),
            self._create_direct_qa(
                "Ku ndodhen bazat ushtarake të Shqipërisë?",
                "Për arsye sigurie, nuk mund të jap detaje specifike për vendndodhjet "
                "e sakta ose kapacitetet e bazave ushtarake. Informacion i përgjithshëm "
                "për strukturat e mbrojtjes mund të gjeni në faqen zyrtare të Ministrisë "
                "së Mbrojtjes (mod.gov.al)."
            ),
        ]

        self.dataset.extend(hardcoded)
        logger.info(f"Added {len(hardcoded)} hardcoded Q&A pairs")

    def build_dataset(self, chunks: List[dict]) -> List[dict]:
        """Build full dataset from chunks."""
        logger.info(f"🚀 Building dataset from {len(chunks)} chunks")

        # 1. Add hardcoded essential knowledge
        self._add_hardcoded_knowledge()

        # 2. Generate Q&A from chunks
        generated = 0
        for chunk in chunks:
            pairs = self._generate_qa_from_chunk(chunk)
            self.dataset.extend(pairs)
            generated += len(pairs)

        logger.info(f"Generated {generated} Q&A pairs from chunks")
        logger.info(f"Total dataset size: {len(self.dataset)}")

        return self.dataset

    def save_dataset(self, train_ratio: float = 0.9):
        """Save dataset as JSONL files with train/validation split."""
        if not self.dataset:
            logger.error("No data to save!")
            return

        # Shuffle
        random.shuffle(self.dataset)

        # Split
        split_idx = int(len(self.dataset) * train_ratio)
        train_data = self.dataset[:split_idx]
        val_data = self.dataset[split_idx:]

        # Save train
        train_path = os.path.join(self.output_dir, "train.jsonl")
        with open(train_path, "w", encoding="utf-8") as f:
            for item in train_data:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")

        # Save validation
        val_path = os.path.join(self.output_dir, "validation.jsonl")
        with open(val_path, "w", encoding="utf-8") as f:
            for item in val_data:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")

        # Save metadata
        metadata = {
            "name": "KIA Dataset",
            "description": "Instruction dataset for Albanian General Staff AI",
            "language": "sq",
            "total_examples": len(self.dataset),
            "train_examples": len(train_data),
            "validation_examples": len(val_data),
            "system_prompt": SYSTEM_PROMPT,
            "format": "ChatML (messages)",
        }
        meta_path = os.path.join(self.output_dir, "metadata.json")
        with open(meta_path, "w", encoding="utf-8") as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)

        logger.info(f"✅ Dataset saved:")
        logger.info(f"   Train: {len(train_data)} examples → {train_path}")
        logger.info(f"   Validation: {len(val_data)} examples → {val_path}")
        logger.info(f"   Metadata: {meta_path}")

    def get_stats(self) -> dict:
        return {
            "total_examples": len(self.dataset),
            "avg_messages_per_example": 3,  # system + user + assistant
        }