import boto3
import hashlib
import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from enum import Enum
from typing import Callable, List, Optional
from langchain_core.embeddings import Embeddings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def default_filename_func(text: str, index: int) -> str:
    md5_hash = hashlib.md5(text.encode('utf-8')).hexdigest()
    return f"{md5_hash}.txt"

class CacheBehavior(Enum):
    NO_CACHE = "no_cache"
    ONLY_CACHE = "only_cache"
    ON_CACHE_MISS_ABORT = "on_cache_miss_abort"
    ON_CACHE_MISS_EMBED_ALL = "on_cache_miss_embed_all"


class S3CachedEmbeddings(Embeddings):
    def __init__(self, 
                 embeddings: Embeddings, 
                 bucket: str, 
                 prefix: Optional[str], 
                 filenaming_function: Optional[Callable[[str, int], str]] = None,
                 cache_behavior = CacheBehavior.NO_CACHE):
        self._s3_client = boto3.client('s3')
        self._bucket = bucket
        self._prefix = prefix
        self._embeddings = embeddings
        self._filenaming_function = filenaming_function
        self._cache_behavior = cache_behavior

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed a list of documents"""


        cache_behavior = self._cache_behavior
        embed_all = cache_behavior == CacheBehavior.NO_CACHE
        if not embed_all:
            logger.info("Using cached embeddings")
            try:
                contents = self._retrieve_content_in_batches(contents=texts, batch_size=len(texts))
                embeddings = [json.loads(content) for content in contents]
            except:
                if cache_behavior == CacheBehavior.ON_CACHE_MISS_EMBED_ALL:
                    logger.info("Cache Miss detected: Embedding all documents")
                    embed_all = True
                else:
                    raise

        if embed_all:
            embeddings = self._embeddings.embed_documents(texts)
            self._upload_content_in_batches(contents=embeddings, batch_size=len(embeddings))

        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """Embed a single text"""
        return self._embeddings.embed_query(text)


    def _upload_content_to_s3(self, content: str, filename: str) -> None:
        """
        Upload content directly to S3 with a generated file name.

        :param bucket_name: S3 bucket name
        :param content: Content to upload
        :param prefix: Prefix for the generated file name
        """
        key = f"{self._prefix}/{filename}"
        try:
            body = json.dumps(content)
            self._s3_client.put_object(Bucket=self._bucket, Key=key, Body=body)
        except Exception as e:
            logger.info(f"Error uploading content: {e}")

    def _upload_content_in_batches(self, contents, batch_size) -> None:
        """
        Upload contents to S3 in batches.

        :param contents: List of strings to upload as content
        :param batch_size: Number of contents to upload concurrently
        :param naming_function: Function to generate file names based on content
        """
        
        naming_function = self._filenaming_function or default_filename_func
        max_workers = min(batch_size, 50)
        for i in range(0, len(contents), batch_size):
            batch = contents[i:i + batch_size]
            logger.info(f"Batch {i // batch_size + 1}: Uploading {len(batch)} embedding files to '{self._bucket}/{self._prefix}'")
            
            # Enumerate the batch to get both content and index
            items_with_index = [(content, i + idx) for idx, content in enumerate(batch)]
            
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                executor.map(
                    lambda item: self._upload_content_to_s3(
                        item[0],  # content
                        naming_function(item[0], item[1]),  # filename generated by naming_function
                    ),
                    items_with_index
                )

    def _retrieve_content_from_s3(self, filename: str) -> str:
        """
        Retrieve content from S3 if the file exists, otherwise return an error message.

        :param filename: The file name to retrieve from S3
        :return: Content of the file or an error message
        """
        key = f"{self._prefix}/{filename}"
        try:
            response = self._s3_client.get_object(Bucket=self._bucket, Key=key)
            content = response['Body'].read().decode('utf-8')
            return content
        except Exception as e:
            logger.error(f"Exception encountered: {e}")
            return e

    
    def _retrieve_content_in_batches(self, contents: list[str], batch_size):
        naming_function = self._filenaming_function or default_filename_func
        max_workers = min(batch_size, 50)
        all_contents = []
        for i in range(0, len(contents), batch_size):
            batch = contents[i:i + batch_size]
            items_with_index = [(content, i + idx) for idx, content in enumerate(batch)]

            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = {
                    executor.submit(self._retrieve_content_from_s3, naming_function(item[0], item[1])): item
                    for item in items_with_index
                }
                

                for future in as_completed(futures):
                    try:
                        result = future.result()
                        if "Error" in result:
                            if self._cache_behavior == CacheBehavior.ONLY_CACHE:
                                logger.warning(f"Warning: Skipping some docs in batch starting at index {i}: {result}")
                                pass
                            else:
                                erm = f"Error processing batch starting at index {i}: {result}"
                                logger.error(erm)
                                raise Exception(erm)
                        else:
                            all_contents.append(result)
                    except Exception as e:
                        if self._cache_behavior == CacheBehavior.ONLY_CACHE:
                            logger.warning(f"Skipping some docs in batch starting at index {i}: {result}")
                            continue
                        else:
                            logger.error(f"Exception encountered: {e}")
                            raise
        return all_contents