import psycopg
import psycopg_pool
import logging
import time
from typing import List, Dict, Any, Tuple
from rag_agent.core.model_client import get_model_client, ModelConfig, ModelType
from rag_agent.core.config import settings
from rag_agent.services.retriever.base_retriever import BaseRetriever
from rag_agent.services.ner_extractor import NERKeywordExtractor
from rag_agent.services.query_expander import QueryExpander

logger = logging.getLogger(__name__)


class TwoStageRetriever(BaseRetriever):
    """
    Two-stage retrieval system:
      - Stage 1: recall-heavy unit filter (documents + sections) using rank fusion (no weights).
      - Stage 2: chunk retrieval with authoritative 70/30 fusion (semantic/keyword).

    Notes:
      - Stage-1 returns a mixed set of document + section UUIDs (deduped by identity).
      - Stage-2 fuses semantic/keyword at the chunk level only (0.7/0.3).
    """

    def __init__(
        self,
        dsn: str,
        top_k: int = 16,
        sql_timeout_s: float = 10.0,
        vector_weight: float = 0.7,
        keyword_weight: float = 0.3,
    ):
        self.dsn = dsn
        # Keep original arg names for compatibility; Stage-1 will internally cap to 75.
        self.top_k = top_k
        self.sql_timeout_s = sql_timeout_s
        self.vector_weight = vector_weight
        self.keyword_weight = keyword_weight

        # Create connection pool for better performance
        self.pool = psycopg_pool.ConnectionPool(
            dsn,
            min_size=1,
            max_size=3,
            timeout=sql_timeout_s
        )

        # Initialize models
        self.embedding_model = settings.EMBEDDING_MODEL
        self.embedding_config = ModelConfig(
            model_type=ModelType.EMBEDDING,
            model_name=self.embedding_model
        )

        # Initialize components
        self.ner_extractor = NERKeywordExtractor()
        self.query_expander = QueryExpander()
        self.embedding_client = get_model_client(self.embedding_config)

    def _embed_query(self, query: str) -> List[float]:
        """Generate embedding for a query."""
        return self.embedding_client.embed_query(query)

    def _log_keyword_only_results_from_data(
        self,
        all_units: List[Dict[str, Any]],
    ) -> None:
        """
        Extract and log keyword-only results from the Stage 1 query data.
        Filters units that have non-zero keyword scores and sorts by keyword rank.
        """
        # Filter units that have keyword matches (non-zero kw_score)
        kw_units = [u for u in all_units if u.get('best_kw_score', 0.0) > 0.0]
        
        # Sort by keyword rank (lower rank = better)
        kw_units_sorted = sorted(kw_units, key=lambda x: x.get('rank_kw', 2147483647))
        
        kw_docs = [u for u in kw_units_sorted if u.get('source_type') == 'document']
        kw_sections = [u for u in kw_units_sorted if u.get('source_type') == 'section']
        
        logger.info(
            f"\nKEYWORD-ONLY RETRIEVAL (Stage 1): {len(kw_docs)} documents, {len(kw_sections)} sections"
        )
        if kw_docs:
            logger.info(f"Keyword Documents (top {min(10, len(kw_docs))}):")
            for i, doc in enumerate(kw_docs[:10], 1):
                logger.info(
                    f"  {i}. {doc['title'][:60]} (kw_score: {doc['best_kw_score']:.4f}, kw_rank: {doc['rank_kw']})"
                )
        if kw_sections:
            logger.info(f"Keyword Sections (top {min(10, len(kw_sections))}):")
            for i, sec in enumerate(kw_sections[:10], 1):
                logger.info(
                    f"  {i}. {sec['title'][:60]} (kw_score: {sec['best_kw_score']:.4f}, kw_rank: {sec['rank_kw']})"
                )
        if not kw_docs and not kw_sections:
            logger.warning("  No keyword-only results found in top 75 fused results!")

    # -----------------------------
    # New Stage-1 (rank fusion) API
    # -----------------------------
    def _stage1_unit_filter_rankfusion(
        self,
        query: str,
        query_vector: List[float],
        include_docs: bool = True,
        include_sections: bool = True,
        cap_units: int = 75,
    ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
        """
        Stage 1: Recall-heavy unit filter with rank fusion (no weights).
        Returns two lists:
          - documents: [{uuid, title, link, best_sem_score, best_kw_score, selector, rank_sem, rank_kw}, ...]
          - sections : [same keys, ...]
        """
        try:
            stage1_sql_path = settings.SQL_DIR / "stage1_document_section_filter.sql"
            stage1_sql = open(stage1_sql_path).read()

            with self.pool.connection() as db_connection:
                with db_connection.cursor() as db_cursor:
                    # Set timeout
                    db_cursor.execute(f"SET LOCAL statement_timeout = {int(self.sql_timeout_s * 1000)};")

                    # Execute Stage 1 (rank fusion) query
                    # NOTE: The SQL itself LIMITs to 75 by default; cap_units is here for clarity if you later parameterize it.
                    db_cursor.execute(
                        stage1_sql,
                        {
                            "query": query,
                            "vector": query_vector,
                            "model_name": self.embedding_model,
                            "include_docs": include_docs,
                            "include_sections": include_sections,
                        }
                    )
                    rows = db_cursor.fetchall()

            documents: List[Dict[str, Any]] = []
            sections: List[Dict[str, Any]] = []
            all_units: List[Dict[str, Any]] = []  # For keyword-only logging

            # Expected columns (by position) from stage1_unit_rank_fusion.sql:
            # 0 source_type, 1 source_uuid, 2 title, 3 link,
            # 4 best_sem_score, 5 best_kw_score, 6 selector, 7 rank_sem, 8 rank_kw
            for r in rows[:cap_units]:
                source_type = r[0]
                unit = {
                    "source_type": source_type,  # Keep for filtering
                    "uuid": r[1],
                    "title": r[2],
                    "link": r[3],
                    "best_sem_score": float(r[4]) if r[4] is not None else 0.0,
                    "best_kw_score": float(r[5]) if r[5] is not None else 0.0,
                    "selector": int(r[6]) if r[6] is not None else 2147483647,
                    "rank_sem": int(r[7]) if r[7] is not None else 2147483647,
                    "rank_kw": int(r[8]) if r[8] is not None else 2147483647,
                }
                all_units.append(unit)  # Add to all_units for keyword logging
                if source_type == "document":
                    documents.append(unit)
                else:
                    sections.append(unit)

            # Log keyword-only results from the data we already fetched
            self._log_keyword_only_results_from_data(all_units)

            logger.info(
                f"\n====================================================== Stage 1: {len(documents)} documents, {len(sections)} sections (total {len(rows)}) ======================================================")
            # Log top-3 per type by selector, then rank_sem
            # for name, items in (("docs", documents), ("secs", sections)):
            #     preview = items[:3]
            #     if preview:
            #         logger.info(
            #             f"   Top {len(preview)} {name}: " +
            #             "; ".join(f"{it['title'][:48]} (sel:{it['selector']}, rS:{it['rank_sem']}, rK:{it['rank_kw']})"
            #                       for it in preview)
            #         )

            return documents, sections

        except Exception as e:
            logger.error(f"Error in Stage 1 (rank fusion) filtering: {e}")
            return [], []

    # --------------------------------------------
    # Back-compat Stage-1 method (kept, minimal)
    # --------------------------------------------
    def _stage1_document_filter(
        self, 
        query: str, 
        query_vector: List[float]
    ) -> List[Dict[str, Any]]:
        """
        Backward-compatible wrapper that returns only documents in the old shape.
        NOTE: Uses the new rank-fusion Stage-1 internally.
        """
        docs, _secs = self._stage1_unit_filter_rankfusion(query, query_vector)
        documents = []
        for doc in docs:
            documents.append({
                'document_uuid': doc['uuid'],
                'title': doc['title'],
                'excerpt': None,  # not provided by rank-fusion Stage-1 (by design)
                # Use best_sem_score as a placeholder for compatibility logging
                'combined_score': float(doc.get('best_sem_score', 0.0))
            })
        logger.info(f"\n====================================================== Stage 1: Returned {len(documents)} documents ======================================================")
        return documents

    # -----------------------------
    # New Stage-2 (70/30 fusion) API
    # -----------------------------
    def _stage2_chunk_retrieval_fusion(
        self,
        query: str,
        query_vector: List[float],
        document_uuids: List[str],
        section_uuids: List[str],
    ) -> List[str]:
        """
        Stage 2: Retrieve chunks with authoritative 70/30 fusion (semantic/keyword).
        Returns formatted chunk blocks:
          <text>...</text>\n<reference><url>...</url></reference>
        """
        try:
            stage2_sql = (
                open(settings.SQL_DIR / "stage2_chunk_retrieval.sql").read()
                .replace("%TOP_K%", str(self.top_k))
            )

            with self.pool.connection() as db_connection:
                with db_connection.cursor() as db_cursor:
                    # Set timeout
                    db_cursor.execute(f"SET LOCAL statement_timeout = {int(self.sql_timeout_s * 1000)};")

                    # Execute Stage 2 (fusion) query
                    db_cursor.execute(
                        stage2_sql,
                        {
                            'query': query,
                            'vector': query_vector,
                            'model_name': self.embedding_model,
                            'document_uuids': document_uuids,
                            'section_uuids': section_uuids,
                            'vector_weight': self.vector_weight,
                            'keyword_weight': self.keyword_weight
                        }
                    )
                    rows = db_cursor.fetchall()

                    # rows: (source_type, source_uuid, content_chunk, link, combined_score)
                    chunks_by_url: Dict[str, str] = {}
                    for row in rows:
                        content_chunk = row[2]
                        link = row[3] if len(row) > 3 and row[3] else ""
                        if link in chunks_by_url:
                            chunks_by_url[link] += f"\n\n{content_chunk}"
                        else:
                            chunks_by_url[link] = content_chunk

                    # Format combined chunks with proper tags for DEFAULT_TEMPLATE
                    chunks = []
                    for link, combined_content in chunks_by_url.items():
                        formatted_chunk = f"<text>{combined_content}</text>"
                        if link:
                            formatted_chunk += f"\n<reference><url>{link}</url></reference>"
                        chunks.append(formatted_chunk)

                    logger.info(f"\n====================================================== Stage 2: Retrieved {len(chunks)} chunk blocks ======================================================")
                    for i, chunk in enumerate(chunks, 1):
                        logger.info(f"\n   {i}. {chunk}{'...' if len(chunk) > 120 else ''}")

                    return chunks

        except Exception as e:
            logger.error(f"Error in Stage 2 chunk retrieval: {e}")
            return []

    # --------------------------------------------
    # Back-compat Stage-2 method (kept, minimal)
    # --------------------------------------------
    def _stage2_chunk_retrieval(
        self, 
        query: str, 
        query_vector: List[float], 
        document_uuids: List[str]
    ) -> List[str]:
        """
        Backward-compatible wrapper that calls the new Stage-2 with only document UUIDs.
        Prefer using _stage2_chunk_retrieval_fusion with both doc & section UUIDs.
        """
        return self._stage2_chunk_retrieval_fusion(
            query=query,
            query_vector=query_vector,
            document_uuids=document_uuids,
            section_uuids=[],
        )

    # --------------
    # Public API
    # --------------
    def retrieve(self, query: str) -> List[str]:
        """
        Two-stage retrieval process:
        1. Extract keywords and enhance query
        2. Stage 1: Rank-fusion unit filter (documents + sections)
        3. Stage 2: Chunk retrieval with authoritative 70/30 fusion
        """
        retrieval_start = time.time()
        logger.info(f"Starting two-stage retrieval for: '{query}'")
        try:
            # Step 1: Embed query
            embed_start = time.time()
            #TODO: Research how to enhance query with NER or query expansion
            # enhanced_query = self.ner_extractor.extract_keywords(query)
            enhanced_query = self.query_expander.expand_query(query)
            logger.info(f"Enhanced query: '{enhanced_query}'")

            query_vector = self._embed_query(query)
            embed_time = time.time() - embed_start

            # Step 2: Stage 1 - Rank fusion (documents + sections)
            stage1_start = time.time()
            doc_units, sec_units = self._stage1_unit_filter_rankfusion(
                enhanced_query, query_vector, include_docs=True, include_sections=True, cap_units=75
            )
            stage1_time = time.time() - stage1_start

            if not doc_units and not sec_units:
                logger.warning("No units found in Stage 1")
                return []

            # Extract UUIDs for Stage 2
            document_uuids = [doc_unit['uuid'] for doc_unit in doc_units]
            section_uuids = [sec_unit['uuid'] for sec_unit in sec_units]

            # Step 3: Stage 2 - Chunk retrieval with 70/30 fusion
            stage2_start = time.time()
            chunks = self._stage2_chunk_retrieval_fusion(
                enhanced_query, query_vector, document_uuids, section_uuids
            )
            stage2_time = time.time() - stage2_start

            # Log retrieval latency
            total_time = time.time() - retrieval_start
            logger.info("")
            logger.info("=" * 60 + " Retrieval Latency " + "=" * 60)
            logger.info(f"Total Retrieval Time: {total_time:.3f}s")
            logger.info(f"  - Embedding:        {embed_time:.3f}s ({embed_time/total_time*100:.1f}%)")
            logger.info(f"  - Stage 1 (Filter): {stage1_time:.3f}s ({stage1_time/total_time*100:.1f}%)")
            logger.info(f"  - Stage 2 (Chunks): {stage2_time:.3f}s ({stage2_time/total_time*100:.1f}%)")
            logger.info(f"  - Other:            {total_time - embed_time - stage1_time - stage2_time:.3f}s")
            logger.info("=" * 140)
            logger.info("")

            return chunks

        except Exception as e:
            logger.error(f"Error in two-stage retrieval: {e}")
            return []

    def __del__(self):
        """Clean up connection pool when object is destroyed."""
        if hasattr(self, 'pool'):
            self.pool.close()
