import json
import time
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from airagagent.pdf_processor import PDFProcessor
from airagagent.vector_store import VectorStore
from airagagent.mistral_integration import RAGAgent
from airagagent.config import PDF_DIR, DOCUMENTS_DIR, DEFAULT_RETRIEVAL_K, MIN_CHUNK_SCORE
from airagagent.exceptions import (
    RAGSystemError,
    VectorStoreError,
    ModelError,
    DocumentProcessingError,
    ValidationError
)
from airagagent.logging_config import get_logger, log_performance, log_error_with_context

class RAGSystem:
    def __init__(self):
        self.logger = get_logger(__name__)
        self.pdf_processor = PDFProcessor()
        self.vector_store = VectorStore()
        self.rag_agent = None  # Initialize later to avoid loading model unnecessarily
        self._initialized = False  # Track initialization state

    def initialize_model(self):
        """Initialize RAG agent (called on demand to save memory)"""
        # No longer loading local model - using Grok API instead
        if self.rag_agent is None:
            self.rag_agent = RAGAgent()  # RAGAgent now uses Grok API, no local model loading

    def setup_system(self):
        """Initialize the RAG system"""
        self.logger.info("RAG System initialization started")

        start_time = time.time()

        try:
            # Load existing vector store
            self.vector_store.load_existing_index()
            self.ensure_enriched_index()

            # Process new PDFs
            new_chunks = self.pdf_processor.process_all_new()

            # Update vector store with new chunks (RAW chunks only - enrichment happens via button)
            if new_chunks:
                self.logger.info(f"Processing {len(new_chunks)} new document chunks")
                self.initialize_model()
                # Use raw chunks directly - enrichment happens manually via UI button
                self.vector_store.add_documents(new_chunks)
                self.vector_store.save_index()
                self.logger.info(f"Successfully processed {len(new_chunks)} raw chunks (enrichment available via UI button)")

            total_docs = len(self.vector_store.documents) if self.vector_store.documents else 0

            setup_duration = time.time() - start_time
            log_performance(self.logger, "system_setup", setup_duration,
                          total_documents=total_docs,
                          new_chunks_processed=len(new_chunks) if new_chunks else 0)

            self.logger.info(f"RAG System initialized successfully with {total_docs} documents")
            return total_docs

        except Exception as e:
            log_error_with_context(self.logger, e, {"operation": "setup_system"}, "Failed to initialize RAG system")
            raise

    def search_and_answer(self, question, k=None):
        """Search for relevant documents and generate answer"""
        start_time = time.time()
        question_preview = str(question)[:100] if question is not None else "None"
        self.logger.info(f"Processing question: {question_preview}{'...' if question and len(str(question)) > 100 else ''}")

        # Use default k if not specified
        if k is None:
            k = DEFAULT_RETRIEVAL_K

        # Validate query early
        if not question:
            self.logger.warning("Empty question provided")
            return {
                'answer': "Please provide a question.",
                'sources': []
            }
        if not isinstance(question, str):
            self.logger.warning(f"Invalid question type: {type(question)}")
            return {
                'answer': "Question must be a string.",
                'sources': []
            }
        question = question.strip()
        if not question:
            self.logger.warning("Question is only whitespace")
            return {
                'answer': "Please provide a meaningful question.",
                'sources': []
            }

        if not self.vector_store.documents:
            self.logger.warning("Search attempted but no documents available")
            return {
                'answer': "No documents available in the system. Please add some PDFs first.",
                'sources': []
            }

        # Initialize model if needed
        try:
            self.initialize_model()
        except Exception as e:
            log_error_with_context(self.logger, e, {"operation": "model_initialization", "question": question[:100]})
            return {
                'answer': f"Error initializing model: {str(e)}",
                'sources': []
            }

        # Detect domain from question for enhanced search
        domain = self._detect_domain_from_question(question)
        self.logger.info(f"Detected domain: {domain}")

        # Use enhanced search with keyword analysis
        try:
            if self._should_use_keyword_focused_search(question):
                # Use keyword-focused search for keyword-heavy queries
                keywords = self.vector_store.keyword_manager.extract_keywords(question, max_keywords=10)
                keyword_list = [kw['keyword'] for kw in keywords]
                search_results = self.vector_store.keyword_focused_search(keyword_list, k=k, domain=domain)
                self.logger.info(f"Used keyword-focused search for {len(keyword_list)} keywords")
            else:
                # Use enhanced general search
                search_results = self.vector_store.enhanced_search(question, k=k, domain=domain)
                self.logger.info("Used enhanced general search")

            self.logger.debug(f"Found {len(search_results)} search results")
            self.logger.debug(f"Found {len(search_results)} search results")
            
            # Log all retrieved chunks with scores for debugging
            self.logger.info(f"Retrieved {len(search_results)} chunks with scores:")
            for i, result in enumerate(search_results[:10], 1):
                score = result.get('score', 0.0)
                vector_score = result.get('vector_score', 0.0)
                keyword_score = result.get('keyword_score', 0.0)
                source = result.get('metadata', {}).get('source', 'Unknown')
                self.logger.info(f"  Chunk {i}: {source[:50]} - Combined: {score:.3f} (Vector: {vector_score:.3f}, Keyword: {keyword_score:.3f})")
            
            # Filter out low-quality chunks below minimum score threshold
            filtered_results = [r for r in search_results if r.get('score', 0.0) >= MIN_CHUNK_SCORE]
            
            if filtered_results:
                self.logger.info(f"Filtered to {len(filtered_results)} high-quality chunks (score >= {MIN_CHUNK_SCORE})")
                score_list = [f"{r.get('score', 0.0):.3f}" for r in filtered_results[:5]]
                self.logger.info(f"Filtered chunk scores: {score_list}")
                search_results = filtered_results
            elif search_results:
                # If all chunks are below threshold, use top 3 anyway but warn
                self.logger.warning(f"All chunks below threshold {MIN_CHUNK_SCORE}, using top 3 anyway")
                top3_scores = [f"{r.get('score', 0.0):.3f}" for r in search_results[:3]]
                self.logger.warning(f"Top 3 scores: {top3_scores}")
                search_results = search_results[:3]
        except ValidationError as e:
            self.logger.warning(f"Query validation failed: {e.message}")
            return {
                'answer': f"Invalid search query: {e.message}",
                'sources': []
            }
        except VectorStoreError as e:
            log_error_with_context(self.logger, e, {"operation": "vector_search", "question": question[:100]})
            return {
                'answer': f"Vector store error: {e.message}",
                'sources': []
            }
        except Exception as e:
            log_error_with_context(self.logger, e, {"operation": "search", "question": question[:100]})
            return {
                'answer': f"Unexpected search error: {str(e)}",
                'sources': []
            }

        if not search_results:
            self.logger.info(f"No relevant documents found for question: {question[:100]}")
            return {
                'answer': "No relevant documents found for your question.",
                'sources': []
            }

        # Format context and generate answer
        try:
            self.logger.debug(f"Formatting context from {len(search_results)} search results")
            context = self.rag_agent.format_context(search_results, question=question)
            self.logger.debug(f"Context length: {len(context)} characters")
            answer = self.rag_agent.generate_answer(context, question, search_results)
            sources = self.rag_agent.prepare_source_summaries(search_results)
            self.logger.info(f"Generated answer with {len(sources)} sources")
            for i, source in enumerate(sources[:5], 1):
                self.logger.info(f"  Source {i}: {source.get('source', 'Unknown')[:50]} - Score: {source.get('score', 0.0):.3f}")

            duration = time.time() - start_time
            log_performance(self.logger, "question_answer", duration,
                          question_length=len(question),
                          results_count=len(search_results),
                          answer_length=len(answer) if answer else 0)

            self.logger.info(f"Successfully answered question in {duration:.2f}s")
            return {
                'answer': answer,
                'sources': sources
            }

        except Exception as e:
            log_error_with_context(self.logger, e, {"operation": "answer_generation", "question": question[:100]})
            return {
                'answer': f"Error generating answer: {str(e)}",
                'sources': []
            }

    def ensure_enriched_index(self):
        """Ensure the vector store uses enriched knowledge cards rather than raw chunks."""
        # OPTIMIZATION: Skip enrichment if vector store already has documents
        # This makes first question much faster - enrichment can happen later if needed
        if self.vector_store.documents and len(self.vector_store.documents) > 0:
            # Check if documents have enriched metadata
            # Safe access with bounds checking
            sample_metadata = {}
            try:
                if len(self.vector_store.documents) > 0:
                    sample_metadata = self.vector_store.documents[0].get('metadata', {})
            except (IndexError, AttributeError) as e:
                self.logger.debug(f"Could not access sample metadata: {e}")
                sample_metadata = {}
            if sample_metadata.get('summary') or sample_metadata.get('card_id'):
                # Already enriched, nothing to do
                return
            else:
                # Has documents but not enriched - skip enrichment for now to speed up initialization
                # The system works fine with raw chunks, enrichment is optional
                self.logger.info(f"Vector store has {len(self.vector_store.documents)} documents (not enriched). Skipping enrichment for faster startup.")
                return
        
        # Only rebuild if vector store is empty
        chunk_files = list(Path(DOCUMENTS_DIR).glob("*_chunks.json"))
        if not chunk_files:
            return

        print("Rebuilding knowledge index from stored chunks...")
        try:
            self._rebuild_vector_store_from_chunks(chunk_files)
        except VectorStoreError as e:
            print(f"Vector store rebuild failed: {e.message}")
            raise
        except Exception as e:
            print(f"Unexpected error during rebuild: {e}")
            raise

    def _rebuild_vector_store_from_chunks(self, chunk_files):
        """Rebuild the vector store using knowledge cards generated from stored chunks."""
        all_chunks = []
        for chunk_file in chunk_files:
            try:
                with open(chunk_file, 'r') as f:
                    data = json.load(f)
                    if isinstance(data, list):
                        all_chunks.extend(data)
            except Exception as e:
                print(f"Warning: Failed to load chunk file {chunk_file}: {e}")

        if not all_chunks:
            print("No chunk data available for rebuilding.")
            return

        print(f"   • Loaded {len(all_chunks)} stored chunks")

        # Reset existing vector store (RAW chunks only - enrichment happens via button)
        try:
            self.vector_store = VectorStore()
            self.initialize_model()
            # Use raw chunks directly - enrichment happens manually via UI button
            if all_chunks:
                self.vector_store.add_documents(all_chunks)
                self.vector_store.save_index()
                print(f"   • Rebuilt vector store with {len(all_chunks)} raw chunks (enrichment available via UI button)")
        except VectorStoreError as e:
            print(f"   • Vector store error: {e.message}")
            raise
        except ModelError as e:
            print(f"   • Model error: {e.message}")
            raise
        except DocumentProcessingError as e:
            print(f"   • Document processing error: {e.message}")
            raise
        except Exception as e:
            print(f"   • Unexpected error rebuilding vector store: {e}")
            raise

    def _detect_domain_from_question(self, question: str) -> Optional[str]:
        """Detect the domain from the question for targeted search."""
        question_lower = question.lower()

        # Domain detection patterns
        domain_patterns = {
            'sports': [
                'football', 'soccer', 'basketball', 'baseball', 'hockey', 'tennis', 'golf',
                'nfl', 'nba', 'mlb', 'nhl', 'premier league', 'la liga', 'bundesliga',
                'odds', 'spread', 'moneyline', 'over/under', 'parlay', 'sports betting'
            ],
            'crypto': [
                'bitcoin', 'btc', 'ethereum', 'eth', 'cryptocurrency', 'crypto', 'blockchain',
                'mining', 'staking', 'defi', 'nft', 'altcoin', 'token', 'wallet', 'exchange'
            ],
            'stocks': [
                'stock', 'equity', 'shares', 'nasdaq', 'nyse', 'dividend', 'earnings',
                'ipo', 'options', 'futures', 'trading', 'investing', 'market cap'
            ],
            'forex': [
                'forex', 'currency', 'fx', 'eur/usd', 'gbp/usd', 'usd/jpy', 'pip',
                'leverage', 'carry trade', 'central bank', 'fed', 'ecb', 'interest rate'
            ]
        }

        # Count matches for each domain
        domain_scores = {}
        for domain, patterns in domain_patterns.items():
            score = sum(1 for pattern in patterns if pattern in question_lower)
            if score > 0:
                domain_scores[domain] = score

        # Return domain with highest score if above threshold
        if domain_scores:
            best_domain = max(domain_scores.items(), key=lambda x: x[1])
            if best_domain[1] >= 1:  # At least one match
                return best_domain[0]

        return None

    def _should_use_keyword_focused_search(self, question: str) -> bool:
        """Determine if query should use keyword-focused search."""
        # Use keyword search for queries with many specific terms
        keywords = self.vector_store.keyword_manager.extract_keywords(question, max_keywords=20)
        high_relevance_keywords = [kw for kw in keywords if kw['score'] > 0.7]

        # Use keyword search if:
        # 1. Many high-relevance keywords found
        # 2. Question contains specific trading terms
        # 3. Question is very specific (longer queries)

        keyword_heavy = len(high_relevance_keywords) >= 3
        specific_terms = any(term in question.lower() for term in [
            'kelly criterion', 'bollinger bands', 'rsi', 'macd', 'fibonacci',
            'parlay', 'round robin', 'arbitrage', 'hedge', 'leverage'
        ])
        long_query = len(question.split()) > 15

        return keyword_heavy or specific_terms or long_query

def main():
    """Command line interface for testing"""
    rag_system = RAGSystem()

    # Setup system
    doc_count = rag_system.setup_system()

    if doc_count == 0:
        print("No documents available. Please add PDFs to the pdf_directory folder.")
        return

    print(f"\nTotal documents in vector store: {doc_count}")

    # Interactive Q&A loop
    while True:
        try:
            question = input("\nEnter your question (or 'quit' to exit): ").strip()

            if question.lower() in ['quit', 'exit', 'q']:
                break

            if not question:
                continue

            print("Searching and generating answer...")
            result = rag_system.search_and_answer(question)

            print(f"\nAnswer: {result['answer']}\n")
            if result['sources']:
                print("Sources:")
                for i, source in enumerate(result['sources'], 1):
                    print(f"{i}. {source['source']} (Score: {source['score']:.3f})")
                    print(f"   Preview: {source['preview']}\n")

        except KeyboardInterrupt:
            print("\nExiting...")
            break
        except Exception as e:
            print(f"Error: {e}")

if __name__ == "__main__":
    main()