"""RAG (Retrieval Augmented Generation) tool implementation."""
from typing import Dict, Any, List
from hawkinsdb import HawkinsDB, LLMInterface
from .base import Tool
import logging
import os
from datetime import datetime

logger = logging.getLogger(__name__)

class RAGTool(Tool):
    """Tool for RAG operations with enhanced file format support"""

    def __init__(self, db_path="knowledge_base.db", chunk_size=500, chunk_overlap=50):
        super().__init__(
            name="rag",
            description="Retrieve and generate content using RAG system"
        )
        logger.info(f"Initializing RAG tool with database: {db_path}")

        # Initialize HawkinsDB and LLM interface
        self.db = HawkinsDB(storage_type='sqlite', db_path=db_path)
        self.llm_interface = LLMInterface(self.db, auto_enrich=True)

        # Configuration for text splitting
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

        logger.info("RAG tool initialized successfully")

    async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """Execute RAG operations"""
        try:
            operation = input_data.get('operation')
            if not operation:
                return {'success': False, 'error': 'No operation specified'}

            logger.info(f"Executing RAG operation: {operation}")

            if operation == 'load':
                return await self._handle_load(input_data)
            elif operation == 'query':
                return await self._handle_query(input_data)
            else:
                return {'success': False, 'error': f'Unknown operation: {operation}'}

        except Exception as e:
            logger.error(f'Error executing RAG operation: {str(e)}')
            return {
                'success': False,
                'error': f'Error executing RAG operation: {str(e)}'
            }

    def _preprocess_content(self, content: str) -> List[str]:
        """Split content into paragraphs while preserving structure"""
        if not content:
            logger.warning("No content provided for preprocessing")
            return []

        logger.debug(f"Raw content length: {len(content)}")
        logger.debug(f"Content preview: {content[:200]}...")

        # Split text into paragraphs
        paragraphs = []
        current_paragraph = []

        for line in content.split('\n'):
            line = line.strip()
            if line:
                current_paragraph.append(line)
            elif current_paragraph:
                paragraph_text = ' '.join(current_paragraph)
                paragraphs.append(paragraph_text)
                logger.debug(f"Created paragraph: {paragraph_text[:100]}...")
                current_paragraph = []

        # Add final paragraph if exists
        if current_paragraph:
            paragraph_text = ' '.join(current_paragraph)
            paragraphs.append(paragraph_text)
            logger.debug(f"Added final paragraph: {paragraph_text[:100]}...")

        logger.info(f"Created {len(paragraphs)} paragraphs from content")
        return paragraphs

    def _create_chunks(self, paragraphs: List[str]) -> List[Dict[str, Any]]:
        """Create chunks from paragraphs with overlap"""
        if not paragraphs:
            logger.warning("No paragraphs provided for chunking")
            return []

        logger.debug(f"Creating chunks from {len(paragraphs)} paragraphs")
        chunks = []
        current_chunk_words = []
        chunk_number = 1

        for paragraph in paragraphs:
            logger.debug(f"Processing paragraph: {paragraph[:50]}...")
            words = paragraph.split()

            # If adding this paragraph would exceed chunk size
            if len(current_chunk_words) + len(words) > self.chunk_size:
                # Save current chunk if it exists
                if current_chunk_words:
                    chunk_text = ' '.join(current_chunk_words)
                    chunks.append({
                        'content': chunk_text,
                        'chunk_number': chunk_number,
                        'length': len(current_chunk_words)
                    })
                    logger.debug(f"Created chunk {chunk_number} with {len(current_chunk_words)} words")
                    chunk_number += 1

                    # Keep overlap words
                    if self.chunk_overlap > 0:
                        current_chunk_words = current_chunk_words[-self.chunk_overlap:]
                    else:
                        current_chunk_words = []

            # Add current paragraph words
            current_chunk_words.extend(words)

        # Add final chunk if it exists
        if current_chunk_words:
            chunk_text = ' '.join(current_chunk_words)
            chunks.append({
                'content': chunk_text,
                'chunk_number': chunk_number,
                'length': len(current_chunk_words)
            })
            logger.debug(f"Created final chunk {chunk_number} with {len(current_chunk_words)} words")

        logger.info(f"Created total of {len(chunks)} chunks")
        return chunks

    async def _handle_load(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle document loading operation"""
        try:
            content = input_data.get('content')
            metadata = input_data.get('metadata', {})

            if not content:
                logger.error("No content provided")
                return {'success': False, 'error': 'No content provided'}

            # Process metadata
            doc_name = metadata.get('name', f'doc_{hash(str(content[:100]))}')
            timestamp = datetime.now().isoformat()
            metadata['timestamp'] = timestamp

            logger.info(f"Processing document: {doc_name}")
            logger.debug(f"Document metadata: {metadata}")

            # Preprocess content
            paragraphs = self._preprocess_content(content)
            if not paragraphs:
                logger.error("No paragraphs created during preprocessing")
                return {'success': False, 'error': 'Failed to create paragraphs from content'}

            # Create chunks
            chunks = self._create_chunks(paragraphs)
            if not chunks:
                logger.error("No chunks created from paragraphs")
                return {'success': False, 'error': 'Failed to create chunks from paragraphs'}

            # Store document entity
            doc_entity = {
                "name": doc_name,
                "column": "Semantic",
                "properties": {
                    "content": content,
                    "total_chunks": len(chunks),
                    "document_length": len(content),
                    "timestamp": timestamp,
                    **metadata
                },
                "relationships": {
                    "has_chunks": []
                }
            }

            # Store chunks with references
            chunk_names = []
            for chunk in chunks:
                chunk_name = f"{doc_name}_chunk_{chunk['chunk_number']}"
                chunk_entity = {
                    "name": chunk_name,
                    "column": "Semantic",
                    "properties": {
                        "content": chunk['content'],
                        "chunk_number": chunk['chunk_number'],
                        "length": chunk['length'],
                        "document": doc_name,
                        **metadata
                    },
                    "relationships": {
                        "part_of": [doc_name]
                    }
                }

                logger.debug(f"Storing chunk {chunk_name} with {chunk['length']} words")
                self.db.add_entity(chunk_entity)
                chunk_names.append(chunk_name)

            # Update document with chunk references and store
            doc_entity["relationships"]["has_chunks"] = chunk_names
            self.db.add_entity(doc_entity)

            logger.info(f"Successfully processed document {doc_name} with {len(chunks)} chunks")
            return {
                'success': True,
                'message': f'Successfully loaded document {doc_name}',
                'chunks_created': len(chunks),
                'doc_name': doc_name
            }

        except Exception as e:
            logger.error(f"Error loading document: {str(e)}")
            return {
                'success': False,
                'error': f'Error loading document: {str(e)}'
            }

    async def _handle_query(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle query operation"""
        try:
            query = input_data.get('query')
            if not query:
                return {'success': False, 'error': 'No query provided'}

            logger.info(f"Processing query: {query}")

            # Enhanced system prompt
            system_prompt = """You are a knowledgeable assistant with access to a document database.
            Your task is to provide accurate and relevant information based on the available context.
            When answering:
            1. Use specific information from the retrieved chunks
            2. Maintain the original meaning and context
            3. If information is not found in the context, clearly state that
            4. Cite relevant sources when available
            5. Be concise and clear in your responses"""

            # Execute query
            response = self.llm_interface.query(
                query,
                system_prompt=system_prompt,
                search_kwargs={
                    "k": 3,  # Number of chunks to retrieve
                    "score_threshold": 0.2,  # Relevance threshold
                    "fetch_k": 10  # Number of candidates for reranking
                }
            )

            # Format response
            if isinstance(response, dict):
                answer = response.get('response', '')
                sources = response.get('source_documents', [])

                formatted_sources = []
                if sources:
                    for source in sources:
                        if isinstance(source, dict):
                            formatted_sources.append({
                                'document': source.get('document', ''),
                                'chunk': source.get('chunk_number', ''),
                                'relevance': source.get('relevance_score', '')
                            })

                return {
                    'success': True,
                    'answer': answer,
                    'sources': formatted_sources
                }
            else:
                return {
                    'success': True,
                    'answer': str(response),
                    'sources': []
                }

        except Exception as e:
            logger.error(f"Error processing query: {str(e)}")
            return {
                'success': False,
                'error': f'Error processing query: {str(e)}'
            }

    def _split_into_paragraphs(self, text: str) -> List[str]:
        """Split text into paragraphs while preserving structure"""
        # This function is no longer used, but keeping it for potential future use.
        # Split on double newlines to preserve paragraph structure
        raw_paragraphs = text.split('\n\n')
        # Clean up paragraphs and remove empty ones
        return [p.strip() for p in raw_paragraphs if p.strip()]