# ============================================================
# 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 re
from pathlib import Path
from typing import TYPE_CHECKING

import pandas as pd

from src.utils.logging import get_logger

if TYPE_CHECKING:
    from src.storage.database import DatabaseManager

logger = get_logger("analytics.panel")


class PanelBuilder:

    def __init__(self, db: DatabaseManager) -> None:
        self.db = db

    def build_panel(
        self,
        tickers: list[str] | None = None,
        form_types: list[str] | None = None,
        start_year: int | None = None,
        end_year: int | None = None,
        frequency: str = "annual",
    ) -> pd.DataFrame:
        conditions: list[str] = []
        params: list = []

        if tickers:
            placeholders = ", ".join(["?"] * len(tickers))
            conditions.append(f"f.ticker IN ({placeholders})")
            params.extend(tickers)
        if form_types:
            placeholders = ", ".join(["?"] * len(form_types))
            conditions.append(f"f.form_type IN ({placeholders})")
            params.extend(form_types)
        if start_year is not None:
            conditions.append("f.fiscal_year >= ?")
            params.append(start_year)
        if end_year is not None:
            conditions.append("f.fiscal_year <= ?")
            params.append(end_year)

        where = f" WHERE {' AND '.join(conditions)}" if conditions else ""

        if frequency == "quarterly":
            group_cols = "f.ticker, c.company_name, f.cik, f.fiscal_year, f.fiscal_quarter, f.form_type"
            select_extra = "f.fiscal_quarter,"
        else:
            group_cols = "f.ticker, c.company_name, f.cik, f.fiscal_year, f.form_type"
            select_extra = ""

        sql = f"""
            SELECT
                f.ticker,
                c.company_name,
                f.cik,
                f.fiscal_year,
                {select_extra}
                f.form_type,
                MAX(f.filing_date) AS filing_date,
                COALESCE(SUM(s.word_count), 0) AS word_count,
                COUNT(DISTINCT s.section_id) AS section_count,
                COALESCE(
                    (SELECT COUNT(*) FROM filing_chunks ch WHERE ch.accession_number = MAX(f.accession_number)),
                    0
                ) AS chunk_count
            FROM filings f
            LEFT JOIN companies c ON f.cik = c.cik
            LEFT JOIN filing_sections s ON f.accession_number = s.accession_number
            {where}
            GROUP BY {group_cols}
            ORDER BY f.ticker, f.fiscal_year
        """

        panel = self.db.fetch_df(sql, params if params else None)
        logger.info("Built panel with %d rows", len(panel))
        return panel

    def add_section_stats(
        self,
        panel: pd.DataFrame,
        section_names: list[str] | None = None,
    ) -> pd.DataFrame:
        if panel.empty:
            return panel

        tickers = panel["ticker"].unique().tolist()
        placeholders = ", ".join(["?"] * len(tickers))

        sql = f"""
            SELECT
                f.ticker,
                f.fiscal_year,
                s.section_name,
                SUM(s.word_count) AS section_word_count
            FROM filing_sections s
            JOIN filings f ON s.accession_number = f.accession_number
            WHERE f.ticker IN ({placeholders})
            GROUP BY f.ticker, f.fiscal_year, s.section_name
        """
        section_df = self.db.fetch_df(sql, tickers)

        if section_df.empty:
            return panel

        if section_names:
            section_df = section_df[section_df["section_name"].isin(section_names)]

        pivoted = section_df.pivot_table(
            index=["ticker", "fiscal_year"],
            columns="section_name",
            values="section_word_count",
            aggfunc="sum",
            fill_value=0,
        )

        pivoted.columns = [self._sanitize_column_name(c) + "_words" for c in pivoted.columns]
        pivoted = pivoted.reset_index()

        result = panel.merge(pivoted, on=["ticker", "fiscal_year"], how="left")
        for col in pivoted.columns:
            if col not in ("ticker", "fiscal_year"):
                result[col] = result[col].fillna(0).astype(int)

        logger.info("Added %d section stat columns", len(pivoted.columns) - 2)
        return result

    def add_custom_measure(
        self,
        panel: pd.DataFrame,
        measure_name: str,
        sql: str,
    ) -> pd.DataFrame:
        measure_df = self.db.fetch_df(sql)

        if measure_df.empty:
            logger.warning("Custom measure '%s' returned no data", measure_name)
            return panel

        merge_cols = [c for c in ["ticker", "fiscal_year", "cik"] if c in measure_df.columns and c in panel.columns]
        if not merge_cols:
            raise ValueError(
                f"Custom measure SQL must return at least one join column "
                f"(ticker, fiscal_year, or cik) matching the panel"
            )

        result = panel.merge(measure_df, on=merge_cols, how="left")
        logger.info("Added custom measure '%s' via columns %s", measure_name, merge_cols)
        return result

    def export_stata(self, panel: pd.DataFrame, output_path: Path) -> None:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        df_clean = panel.copy()
        for col in df_clean.columns:
            if df_clean[col].dtype == "object":
                df_clean[col] = df_clean[col].astype(str).str[:244]

        df_clean.columns = [self._sanitize_column_name(c) for c in df_clean.columns]

        df_clean.to_stata(str(output_path), write_index=False, version=118)
        logger.info("Exported panel (%d rows) to Stata: %s", len(df_clean), output_path)

    def export_csv(self, panel: pd.DataFrame, output_path: Path) -> None:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        panel.to_csv(output_path, index=False)
        logger.info("Exported panel (%d rows) to CSV: %s", len(panel), output_path)

    @staticmethod
    def _sanitize_column_name(name: str) -> str:
        sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", name.lower().strip())
        sanitized = re.sub(r"_+", "_", sanitized).strip("_")
        if sanitized and sanitized[0].isdigit():
            sanitized = "col_" + sanitized
        return sanitized[:32]