# ============================================================
# 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 datetime import datetime, timezone
from pathlib import Path
from typing import Any

import duckdb
import pandas as pd
import polars as pl

from src.storage.schema import DatabaseSchema
from src.utils.constants import DUCKDB_PATH
from src.utils.logging import get_logger

logger = get_logger("storage.database")


class DatabaseManager:

    def __init__(self, db_path: Path | str | None = None) -> None:
        self.db_path = Path(db_path) if db_path else DUCKDB_PATH
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = duckdb.connect(str(self.db_path))
        logger.info("Connected to DuckDB at %s", self.db_path)

    def __enter__(self) -> DatabaseManager:
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.close()

    def close(self) -> None:
        if self.conn:
            self.conn.close()
            self.conn = None
            logger.info("Closed DuckDB connection")

    def initialize(self) -> None:
        for sql in DatabaseSchema.get_create_tables_sql():
            self.conn.execute(sql)
        for sql in DatabaseSchema.get_upgrade_sql():
            try:
                self.conn.execute(sql)
            except Exception:  # noqa: BLE001 - column already present
                pass
        for sql in DatabaseSchema.get_create_indexes_sql():
            self.conn.execute(sql)
        for sql in DatabaseSchema.get_create_views_sql():
            self.conn.execute(sql)
        logger.info("Database schema initialized")

    def execute(self, sql: str, params: list | tuple | None = None) -> duckdb.DuckDBPyRelation:
        if params:
            return self.conn.execute(sql, params)
        return self.conn.execute(sql)

    def execute_many(self, sql: str, data: list) -> None:
        self.conn.executemany(sql, data)

    def fetch_df(self, sql: str, params: list | tuple | None = None) -> pd.DataFrame:
        result = self.execute(sql, params)
        return result.fetchdf()

    def fetch_polars(self, sql: str, params: list | tuple | None = None) -> pl.DataFrame:
        result = self.execute(sql, params)
        return result.pl()

    # ---- CRUD: companies ----

    def insert_company(self, company: dict) -> None:
        cols = list(company.keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO companies ({col_names}) VALUES ({placeholders})"
        self.conn.execute(sql, list(company.values()))

    def insert_companies(self, companies: list[dict]) -> None:
        if not companies:
            return
        cols = list(companies[0].keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO companies ({col_names}) VALUES ({placeholders})"
        self.conn.executemany(sql, [list(c.values()) for c in companies])
        logger.info("Inserted %d companies", len(companies))

    # ---- CRUD: filings ----

    def insert_filing(self, filing: dict) -> None:
        cols = list(filing.keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO filings ({col_names}) VALUES ({placeholders})"
        self.conn.execute(sql, list(filing.values()))

    def insert_filings(self, filings: list[dict]) -> None:
        if not filings:
            return
        cols = list(filings[0].keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO filings ({col_names}) VALUES ({placeholders})"
        self.conn.executemany(sql, [list(f.values()) for f in filings])
        logger.info("Inserted %d filings", len(filings))

    # ---- CRUD: sections ----

    def insert_sections(self, sections: list[dict]) -> None:
        if not sections:
            return
        cols = list(sections[0].keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO filing_sections ({col_names}) VALUES ({placeholders})"
        self.conn.executemany(sql, [list(s.values()) for s in sections])
        logger.info("Inserted %d sections", len(sections))

    # ---- CRUD: chunks ----

    def insert_chunks(self, chunks: list[dict]) -> None:
        if not chunks:
            return
        cols = list(chunks[0].keys())
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        sql = f"INSERT OR REPLACE INTO filing_chunks ({col_names}) VALUES ({placeholders})"
        self.conn.executemany(sql, [list(c.values()) for c in chunks])
        logger.info("Inserted %d chunks", len(chunks))

    _PRIMARY_KEYS = {
        "companies": "cik", "filings": "accession_number", "filing_sections": "section_id",
        "filing_chunks": "chunk_id", "embeddings_metadata": "embedding_id",
    }

    # columns that are indexed / part of a foreign key: DuckDB rewrites their UPDATE as delete+insert,
    # which fails on referenced rows, so they are never updated in place
    _FROZEN_COLS = {"filings": {"cik", "form_type", "filing_date", "fiscal_year"}}

    def insert_df(self, table: str, df: pd.DataFrame, upsert: bool = True) -> int:
        """Bulk upsert from a DataFrame (orders of magnitude faster than executemany).

        Existing rows are UPDATEd column by column and new rows INSERTed, which is
        the only form DuckDB accepts for rows referenced by foreign keys
        (companies <- filings <- sections <- chunks).
        """
        if df is None or df.empty:
            return 0
        cols = list(df.columns)
        col_list = ", ".join(cols)
        pk = self._PRIMARY_KEYS.get(table)
        self.conn.register("_df_in", df)
        try:
            if upsert and pk and pk in cols:
                frozen = self._FROZEN_COLS.get(table, set())
                upd_cols = [c for c in cols if c != pk and c not in frozen]
                if upd_cols:
                    updates = ", ".join(f"{c} = d.{c}" for c in upd_cols)
                    self.conn.execute(
                        f"UPDATE {table} SET {updates} FROM _df_in d WHERE {table}.{pk} = d.{pk}"
                    )
                self.conn.execute(
                    f"INSERT INTO {table} ({col_list}) SELECT {col_list} FROM _df_in d "
                    f"WHERE NOT EXISTS (SELECT 1 FROM {table} t WHERE t.{pk} = d.{pk})"
                )
            else:
                self.conn.execute(f"INSERT OR REPLACE INTO {table} ({col_list}) SELECT {col_list} FROM _df_in")
        finally:
            self.conn.unregister("_df_in")
        return len(df)

    def insert_records_df(self, table: str, records: list[dict]) -> int:
        if not records:
            return 0
        return self.insert_df(table, pd.DataFrame.from_records(records))

    # ---- Status updates ----

    def update_filing_status(self, accession_number: str, field: str, status: str) -> None:
        allowed_fields = {"download_status", "parse_status"}
        if field not in allowed_fields:
            raise ValueError(f"field must be one of {allowed_fields}, got '{field}'")
        sql = f"UPDATE filings SET {field} = ?, updated_at = current_timestamp WHERE accession_number = ?"
        self.conn.execute(sql, [status, accession_number])

    def log_operation(
        self,
        operation: str,
        entity_type: str,
        entity_id: str,
        status: str,
        message: str = "",
    ) -> None:
        log_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        sql = """
            INSERT INTO processing_log (log_id, operation, entity_type, entity_id, status, message, started_at, completed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """
        self.conn.execute(sql, [log_id, operation, entity_type, entity_id, status, message, now, now])

    # ---- Query methods ----

    def get_filings(
        self,
        cik: str | None = None,
        ticker: str | None = None,
        form_type: str | None = None,
        start_date: str | None = None,
        end_date: str | None = None,
        fiscal_year: int | None = None,
    ) -> pd.DataFrame:
        conditions: list[str] = []
        params: list[Any] = []

        if cik:
            conditions.append("cik = ?")
            params.append(cik)
        if ticker:
            conditions.append("ticker = ?")
            params.append(ticker)
        if form_type:
            conditions.append("form_type = ?")
            params.append(form_type)
        if start_date:
            conditions.append("filing_date >= ?")
            params.append(start_date)
        if end_date:
            conditions.append("filing_date <= ?")
            params.append(end_date)
        if fiscal_year is not None:
            conditions.append("fiscal_year = ?")
            params.append(fiscal_year)

        where = f" WHERE {' AND '.join(conditions)}" if conditions else ""
        sql = f"SELECT * FROM filings{where} ORDER BY filing_date DESC"
        return self.fetch_df(sql, params if params else None)

    def get_sections(self, accession_number: str) -> pd.DataFrame:
        sql = "SELECT * FROM filing_sections WHERE accession_number = ? ORDER BY section_order"
        return self.fetch_df(sql, [accession_number])

    def get_chunks(
        self,
        accession_number: str | None = None,
        section_id: str | None = None,
    ) -> pd.DataFrame:
        conditions: list[str] = []
        params: list[Any] = []

        if accession_number:
            conditions.append("accession_number = ?")
            params.append(accession_number)
        if section_id:
            conditions.append("section_id = ?")
            params.append(section_id)

        where = f" WHERE {' AND '.join(conditions)}" if conditions else ""
        sql = f"SELECT * FROM filing_chunks{where} ORDER BY chunk_order"
        return self.fetch_df(sql, params if params else None)

    def get_companies(self) -> pd.DataFrame:
        return self.fetch_df("SELECT * FROM companies ORDER BY ticker")

    def get_filing_count(self) -> int:
        result = self.conn.execute("SELECT COUNT(*) FROM filings").fetchone()
        return result[0]

    def get_stats(self) -> dict:
        tables = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata", "processing_log"]
        stats: dict[str, int] = {}
        for table in tables:
            result = self.conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()
            stats[table] = result[0]
        return stats