# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

from __future__ import annotations

import uuid
from pathlib import Path
from typing import TYPE_CHECKING

import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq
import pandas as pd

from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL
from src.utils.logging import get_logger

if TYPE_CHECKING:
    from src.storage.database import DatabaseManager

logger = get_logger("storage.parquet")


class ParquetManager:

    def __init__(self, base_dir: Path | str | None = None) -> None:
        self.base_dir = Path(base_dir) if base_dir else DATA_DIRS["parquet"]
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def _table_dir(self, table_name: str) -> Path:
        path = self.base_dir / table_name
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _write_dataset(
        self,
        table: pa.Table,
        table_name: str,
        partition_by: list[str] | None = None,
    ) -> None:
        out_dir = self._table_dir(table_name)
        pq.write_to_dataset(
            table,
            root_path=str(out_dir),
            partition_cols=partition_by,
            compression=PARQUET_COMPRESSION,
            compression_level=PARQUET_COMPRESSION_LEVEL,
            existing_data_behavior="overwrite_or_ignore",
        )
        logger.info(
            "Wrote %d rows to %s (partitions=%s)",
            table.num_rows,
            out_dir,
            partition_by,
        )

    def write_filings(
        self,
        filings: list[dict],
        partition_by: list[str] | None = None,
    ) -> None:
        if partition_by is None:
            partition_by = ["form_type", "fiscal_year"]
        table = pa.Table.from_pylist(filings)
        self._write_dataset(table, "filings", partition_by)

    def write_sections(
        self,
        sections: list[dict],
        partition_by: list[str] | None = None,
    ) -> None:
        if partition_by is None:
            partition_by = ["form_type", "fiscal_year"]
        table = pa.Table.from_pylist(sections)
        self._write_dataset(table, "filing_sections", partition_by)

    def write_chunks(
        self,
        chunks: list[dict],
        partition_by: list[str] | None = None,
    ) -> None:
        if partition_by is None:
            partition_by = ["form_type", "fiscal_year"]
        table = pa.Table.from_pylist(chunks)
        self._write_dataset(table, "filing_chunks", partition_by)

    def write_embeddings(self, embeddings: list[dict], model_name: str) -> None:
        for emb in embeddings:
            emb.setdefault("model_name", model_name)
        table = pa.Table.from_pylist(embeddings)
        self._write_dataset(table, "embeddings", partition_by=["model_name"])

    def write_embeddings_table(self, table: pa.Table, model_name: str) -> Path:
        """Write one parquet file per call (no tiny-file explosion for large corpora)."""
        out_dir = self._table_dir("embeddings") / f"model_name={model_name}"
        out_dir.mkdir(parents=True, exist_ok=True)
        path = out_dir / f"emb-{uuid.uuid4().hex}.parquet"
        pq.write_table(table.drop_columns(["model_name"]) if "model_name" in table.column_names else table,
                       str(path), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL)
        return path

    def read_table(self, table_name: str, filters: list | None = None) -> pa.Table:
        table_dir = self._table_dir(table_name)
        dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")
        if filters:
            return dataset.to_table(filter=ds.field(*filters) if len(filters) == 3 else None)
        return dataset.to_table()

    def read_df(self, table_name: str, filters: list | None = None) -> pd.DataFrame:
        return self.read_table(table_name, filters).to_pandas()

    def export_to_parquet(self, db: DatabaseManager, table_name: str) -> None:
        df = db.fetch_df(f"SELECT * FROM {table_name}")
        table = pa.Table.from_pandas(df)
        out_dir = self._table_dir(table_name)
        pq.write_table(
            table,
            str(out_dir / f"{table_name}.parquet"),
            compression=PARQUET_COMPRESSION,
            compression_level=PARQUET_COMPRESSION_LEVEL,
        )
        logger.info("Exported table '%s' (%d rows) to parquet", table_name, len(df))

    def import_from_parquet(self, db: DatabaseManager, table_name: str) -> None:
        table_dir = self._table_dir(table_name)
        dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")
        arrow_table = dataset.to_table()
        df = arrow_table.to_pandas()
        db.conn.execute(f"INSERT OR REPLACE INTO {table_name} SELECT * FROM df")
        logger.info("Imported %d rows from parquet into table '%s'", len(df), table_name)

    def get_dataset_info(self, table_name: str) -> dict:
        table_dir = self._table_dir(table_name)
        parquet_files = list(table_dir.rglob("*.parquet"))
        total_size = sum(f.stat().st_size for f in parquet_files)

        row_count = 0
        if parquet_files:
            dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")
            row_count = dataset.count_rows()

        return {
            "table_name": table_name,
            "row_count": row_count,
            "file_count": len(parquet_files),
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "path": str(table_dir),
        }