# ============================================================
# 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

from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

from src.utils.logging import get_logger

if TYPE_CHECKING:
    from src.embeddings.manager import EmbeddingManager
    from src.storage.database import DatabaseManager

logger = get_logger("search.semantic")


class SemanticSearch:

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

    def search(
        self,
        query: str,
        top_k: int = 20,
        form_type: str | None = None,
        ticker: str | None = None,
        start_date: str | None = None,
        end_date: str | None = None,
        source_type: str = "chunk",
    ) -> pd.DataFrame:
        model_name = self.emb.model_name
        query_vector = self.emb.registry.encode([query], model_name=model_name, show_progress=False)
        query_vector = query_vector[0]

        conditions: list[str] = ["em.source_type = ?"]
        params: list = [source_type]

        if form_type:
            conditions.append("f.form_type = ?")
            params.append(form_type)
        if ticker:
            conditions.append("f.ticker = ?")
            params.append(ticker)
        if start_date:
            conditions.append("f.filing_date >= ?")
            params.append(start_date)
        if end_date:
            conditions.append("f.filing_date <= ?")
            params.append(end_date)

        conditions.append("em.model_name = ?")
        params.append(model_name)

        where_clause = " AND ".join(conditions)

        if source_type == "chunk":
            sql = f"""
                SELECT
                    em.source_id,
                    c.chunk_text AS text,
                    c.accession_number,
                    f.ticker,
                    f.company_name,
                    f.filing_date,
                    f.form_type,
                    s.section_name
                FROM embeddings_metadata em
                JOIN filing_chunks c ON em.source_id = c.chunk_id
                JOIN filings f ON c.accession_number = f.accession_number
                LEFT JOIN filing_sections s ON c.section_id = s.section_id
                WHERE {where_clause}
            """
        else:
            sql = f"""
                SELECT
                    em.source_id,
                    sec.text,
                    sec.accession_number,
                    f.ticker,
                    f.company_name,
                    f.filing_date,
                    f.form_type,
                    sec.section_name
                FROM embeddings_metadata em
                JOIN filing_sections sec ON em.source_id = sec.section_id
                JOIN filings f ON sec.accession_number = f.accession_number
                WHERE {where_clause}
            """

        candidates_df = self.db.fetch_df(sql, params)

        if candidates_df.empty:
            logger.info("No candidates found for query")
            return pd.DataFrame()

        source_ids = candidates_df["source_id"].tolist()
        vectors = self.emb.get_vectors(source_ids, model_name=model_name)

        if vectors.size == 0:
            logger.warning("No embedding vectors found for candidates")
            return pd.DataFrame()

        similarities = self._cosine_similarity(query_vector, vectors)
        candidates_df = candidates_df.copy()
        candidates_df["similarity_score"] = similarities

        results = (
            candidates_df
            .sort_values("similarity_score", ascending=False)
            .head(top_k)
            .reset_index(drop=True)
        )

        logger.info(
            "Search returned %d results (top score=%.4f)",
            len(results),
            results["similarity_score"].iloc[0] if not results.empty else 0.0,
        )
        return results

    def search_similar(self, source_id: str, top_k: int = 10) -> pd.DataFrame:
        meta_sql = "SELECT source_type, model_name FROM embeddings_metadata WHERE source_id = ? LIMIT 1"
        meta_df = self.db.fetch_df(meta_sql, [source_id])

        if meta_df.empty:
            logger.warning("Source ID '%s' not found in embeddings metadata", source_id)
            return pd.DataFrame()

        source_type = meta_df["source_type"].iloc[0]
        model_name = meta_df["model_name"].iloc[0]

        source_vector = self.emb.get_vectors([source_id], model_name=model_name)
        if source_vector.size == 0:
            return pd.DataFrame()
        source_vector = source_vector[0]

        all_meta_sql = (
            "SELECT source_id FROM embeddings_metadata "
            "WHERE source_type = ? AND model_name = ? AND source_id != ?"
        )
        all_meta_df = self.db.fetch_df(all_meta_sql, [source_type, model_name, source_id])

        if all_meta_df.empty:
            return pd.DataFrame()

        candidate_ids = all_meta_df["source_id"].tolist()
        vectors = self.emb.get_vectors(candidate_ids, model_name=model_name)

        if vectors.size == 0:
            return pd.DataFrame()

        similarities = self._cosine_similarity(source_vector, vectors)

        if source_type == "chunk":
            detail_sql = """
                SELECT c.chunk_id AS source_id, c.chunk_text AS text, c.accession_number,
                       f.ticker, f.company_name, f.filing_date, f.form_type
                FROM filing_chunks c
                JOIN filings f ON c.accession_number = f.accession_number
                WHERE c.chunk_id IN ({})
            """.format(", ".join(["?"] * len(candidate_ids)))
        else:
            detail_sql = """
                SELECT s.section_id AS source_id, s.text, s.accession_number,
                       f.ticker, f.company_name, f.filing_date, f.form_type
                FROM filing_sections s
                JOIN filings f ON s.accession_number = f.accession_number
                WHERE s.section_id IN ({})
            """.format(", ".join(["?"] * len(candidate_ids)))

        detail_df = self.db.fetch_df(detail_sql, candidate_ids)

        sim_df = pd.DataFrame({"source_id": candidate_ids, "similarity_score": similarities})
        results = detail_df.merge(sim_df, on="source_id")
        results = (
            results
            .sort_values("similarity_score", ascending=False)
            .head(top_k)
            .reset_index(drop=True)
        )
        return results

    def search_companies(self, query: str, top_k: int = 5) -> pd.DataFrame:
        results = self.search(query, top_k=200)

        if results.empty:
            return pd.DataFrame()

        company_agg = (
            results
            .groupby(["ticker", "company_name"])
            .agg(
                avg_similarity=("similarity_score", "mean"),
                max_similarity=("similarity_score", "max"),
                match_count=("similarity_score", "count"),
            )
            .reset_index()
            .sort_values("avg_similarity", ascending=False)
            .head(top_k)
            .reset_index(drop=True)
        )
        return company_agg

    @staticmethod
    def _cosine_similarity(query_vec: np.ndarray, matrix: np.ndarray) -> np.ndarray:
        query_norm = np.linalg.norm(query_vec)
        if query_norm == 0:
            return np.zeros(matrix.shape[0])
        matrix_norms = np.linalg.norm(matrix, axis=1)
        matrix_norms[matrix_norms == 0] = 1.0
        return np.dot(matrix, query_vec) / (matrix_norms * query_norm)