import pandas as pd
from tqdm import tqdm
from src.datasets.utils import word_count
from src.logger import get_logger

logger = get_logger(__name__)


def extract_text(files: list[str], lang: str = "pl", min_wc: int = 0) -> list[str]:
    logger.info(f"Extracting text data for language: {lang}")

    texts = set()
    for filename in tqdm(files, desc="Processing files"):
        chunk_df = pd.read_csv(filename, sep="\t")

        chunk_filter = (chunk_df["language"] == lang) & (
            ~chunk_df["context_page_description"].isna()
            | ~chunk_df["context_section_description"].isna()
        )
        chunk_df = chunk_df[chunk_filter]

        length_filter = lambda df, col: df[col].apply(word_count) >= min_wc
        texts.update(chunk_df[length_filter(chunk_df, "context_page_description")]["context_page_description"].dropna().tolist())
        texts.update(chunk_df[length_filter(chunk_df, "context_section_description")]["context_section_description"].dropna().tolist())

    logger.info(f"Extracted {len(texts)} unique text entries (with at least {min_wc} words)")
    return list(texts)


def save_text(texts: list[str], output_path: str):
    with open(output_path, "w") as f:
        for text in texts:
            f.write(text + "\n")

if __name__ == "__main__":
    import os
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="Path to the directory containing raw WIT files",
    )
    parser.add_argument(
        "--output_path",
        type=str,
        required=True,
        help="Path to save the extracted text data",
    )
    parser.add_argument("--lang", type=str, default="pl", help="Language code to filter the data")
    parser.add_argument("--min_wc", type=int, default=0, help="Minimal word count of single entry")
    args = parser.parse_args()

    datafiles = [
        os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if f.endswith(".tsv")
    ]
    logger.info(f"Found {len(datafiles)} input files.")

    texts = extract_text(datafiles, lang=args.lang, min_wc=args.min_wc)

    logger.info(f"Saving extracted text data to {args.output_path}")
    save_text(texts, args.output_path)
