import sqlite3
import pandas as pd
import os
import itertools
import ast
import uuid
from huggingface_hub import hf_hub_download, upload_file
from pathlib import Path

# Initialisations
print("Initialisations ...")
hf_token = os.environ["API_HF_TOKEN"]
DATA_DIR = Path("../data")   # dossier parent du script
parquet_path = hf_hub_download(repo_id="Loren/articles_data",
                               filename="mon_fichier.parquet")
REPO_ID = "Loren/articles_db"  # dataset HF
DB_NAME = 'articles.db'
SQLITE_FILE = DATA_DIR / DB_NAME
# Créer le dossier data s'il n'existe pas
DATA_DIR.mkdir(exist_ok=True)

# Chargement des données
print("Chargement des données ...")
df = pd.read_parquet(parquet_path)

# Initialisations de la base SQLite
print("Initialisations de la base SQLite ...")
conn = sqlite3.connect(SQLITE_FILE)
cur = conn.cursor()

# Suppression des anciennes tables
cur.execute("DROP TABLE IF EXISTS tag_article")
cur.execute("DROP TABLE IF EXISTS tags")
cur.execute("DROP TABLE IF EXISTS articles")

# Création des tables Articles, Tags, et de la table d'association articles <-> tags
cur.execute("""
CREATE TABLE articles (
    article_id TEXT PRIMARY KEY,         -- UUID
    article_title TEXT,
    article_text TEXT,
    article_url TEXT,
    article_authors TEXT,
    article_date TEXT                    -- YYYY-MM-DD
)""")

cur.execute("""
CREATE TABLE tags (
    tag_id INTEGER PRIMARY KEY AUTOINCREMENT,
    tag_name TEXT UNIQUE
)""")

cur.execute("""
CREATE TABLE tag_article (
    tag_article_id INTEGER PRIMARY KEY AUTOINCREMENT,
    article_id TEXT,
    tag_id INTEGER,
    FOREIGN KEY(article_id) REFERENCES articles(article_id),
    FOREIGN KEY(tag_id) REFERENCES tags(tag_id)
)""")

# Extraction des tags en une liste unique
print("Extraction des tags en une liste unique ...")
df['list_tags'] = df['tags'].apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else [])
# Extraire tous les tags uniques
all_tags = list(set(itertools.chain.from_iterable(df['list_tags'])))

# Insertion des tags dans la table
print("Insertion des tags dans la table ...")
cur.executemany("INSERT INTO tags (tag_name) VALUES (?)", [(tag,) for tag in all_tags])

# Récupération des correspondances tag_name -> tag_id
print("Récupération des correspondances tag_name -> tag_id ...")
cur.execute("SELECT tag_id, tag_name FROM tags")
dict_tag_map = {tag_name: tag_id for tag_id, tag_name in cur.fetchall()}

# Insertion des articles et table d'association dans les tables
print("Insertion des articles et table d'association dans les tables ...")
for _, row in df.iterrows():
    # Détermination de l'id article
    article_id = str(uuid.uuid4())

    # Extraction de la date du timestamp
    date_value = None
    if pd.notna(row["timestamp"]):
        try:
            date_value = str(pd.to_datetime(row["timestamp"]).date())
        except Exception:
            date_value = None

    # Insertion dans la table Articles
    cur.execute("""
        INSERT INTO articles (article_id, article_title, article_text, article_url, article_authors, article_date)
        VALUES (?, ?, ?, ?, ?, ?)""",
               (article_id, row["title"], row["text"], row["url"], row["authors"], date_value))

    # Association aux tags
    for tag_name in row['list_tags']:
        tag_id = dict_tag_map[tag_name]
        cur.execute("INSERT INTO tag_article (article_id, tag_id) VALUES (?, ?)",
                    (article_id, tag_id))

conn.commit()
conn.close()

# Upload dans le dataset hugging face
print("Upload dans le dataset hugging face ...")
upload_file(
    path_or_fileobj=SQLITE_FILE,
    path_in_repo=DB_NAME,
    repo_id=REPO_ID,
    repo_type="dataset",
    token=hf_token
)

print("Traitement terminé.")