import logging

import polars as pl

from utils.embed import embed as embed
from utils.paths import DATA

logger = logging.getLogger(__name__)

logging.basicConfig(level=logging.INFO)


def load_dataset(file_name: str):
    features = ["meta_title", "meta_description", "content"]
    return (
        pl.scan_csv(file_name)
        .with_columns(
            pl.concat_str([pl.col(c) for c in features], separator="\n\n").alias(
                "text"
            ),
            pl.col("date").str.to_date().alias("date"),
        )
        .rename(
            {
                "is_news_article": "is_news",
                "link_count": "links",
                "paragraph_count": "paragraphs",
            }
        )
        .select("text", "is_news", "url", "date", "paragraphs", "links")
        .collect()
    )


def main() -> None:
    for name in ["train", "eval"]:
        df = load_dataset(DATA / (name + ".csv"))
        embeds = embed(df.get_column("text").to_list())
        df = df.with_columns(pl.Series(embeds).alias("embeds")).write_parquet(
            DATA / (name + ".parquet")
        )


if __name__ == "__main__":
    main()