import os, asyncio
from abc import ABC, abstractmethod
from typing import List, Dict, Tuple

from aiobotocore.session import get_session

from qdrant_client import models

from animuz_core.embedding.embedding_client import EmbeddingClient
from animuz_core.embedding.s3_embedding_client import S3EmbeddingClient
from animuz_core.ingest.custom_unstructured_client import MyUnstructuredClient
from animuz_core.ingest.azure_doc_ai_client import AzureDocAiClient
from animuz_core.vectordb.utils import MyQdrantClient
from animuz_core.ingest.structured import Structured
    
class ResponseGetChunk:
    text: str
    source: str

def get_file_extension(key: str) -> str:
    parts = key.split('.')
    return parts[-1]

class BasePipeline(ABC):
    def __init__(self, embedding_client: EmbeddingClient, 
                 db_client: MyQdrantClient, 
                 unstructured_client: MyUnstructuredClient = None, 
                 azure_doc_ai_client: AzureDocAiClient = None
                 ) -> None:
        
        self.embedding_client = embedding_client
        self.db_client = db_client

        self.unstructured_client = unstructured_client if unstructured_client else MyUnstructuredClient(host=os.environ.get("UNSTRUCTURED_HOST"), port=os.environ.get("UNSTRUCTURED_PORT"))
        self.azure_doc_ai_client = azure_doc_ai_client if azure_doc_ai_client else AzureDocAiClient(endpoint=os.environ.get("AZURE_DOC_AI_ENDPOINT"), key=os.environ.get("AZURE_DOC_AI_KEY"))
        
    async def parse_doc(self, path: str) -> List[dict]:

        assert os.path.isfile(path), "Must be a valid file"

        if not path.endswith(".pdf"):
            res = await self.unstructured_client.a_partition(path)
        else:
            res = await self.azure_doc_ai_client.analyze_document(path) #TODO: maybe find an open source for this

        return res
    
    async def parse_doc_markdown(self, path: str, is_chunked: bool) -> Tuple[List[ResponseGetChunk], int]:
        print('========== START PARSE DOC MARKDOWN ========== ')
        print('Path => ', path)
        assert os.path.isfile(path), "Must be a valid file"

        with open(path, "rb") as f:
            buffer = f.read()

        file_extension = get_file_extension(path)

        if is_chunked:
            return Structured.get_chunks(buffer)

        if file_extension in ['txt', 'md', 'tex', 'csv']:
            return Structured.split_file(buffer, file_extension)
        elif file_extension in ['pdf']:
            # The charcount of pdf will be chunked counts instead of file counts (but shouldn't be used)
            res = await self.azure_doc_ai_client.analyze_document(path)
            return res, sum([len(chunk.text) for chunk in res])
        else:
            # The charcount of pdf will be chunked counts instead of file counts (but shouldn't be used)
            res = await self.unstructured_client.a_partition(path)
            return res, sum([len(chunk.text) for chunk in res])


    async def batch_embed(self, texts: List[str]) -> List[dict]:
        return await self.embedding_client.batch_get_embedding(texts)
    
    async def upload_to_db(self, dense_embeddings: List[List[float]], sparse_indices: List[List[int]], sparse_values: List[List[float]], properties: List[Dict[str, str]]) -> None:
        await self.db_client.upload_batch(dense_embeddings, sparse_indices, sparse_values, properties)
    
    async def upload_to_db_v2(self, dense_embeddings: List[List[float]], list_of_sparse_vectors: List[Dict[str, Dict[str, List]]], properties: List[Dict[str, str]]) -> None:
        await self.db_client.upload_batch_v2(dense_embeddings, list_of_sparse_vectors, properties)

    async def embed_single(self, query: str) -> Dict[str, float]:
        return await self.embedding_client.get_embedding(query)
    
    async def query_db(self, query: str, user_chat_id: str, hybrid: bool = True, top_k: int = 3) -> Tuple[List[str], List[models.ScoredPoint]]:
        emb_res = await self.embed_single(query)

        indices = emb_res["sparse_embedding"].keys()
        values = emb_res["sparse_embedding"].values()

        if hybrid:
            search_results = await self.db_client.hybrid_search(emb_res["embedding"], indices, values, user_chat_id, top_k=top_k)
            return [res["payload"]["text"] for res in search_results], search_results
        else:
            raise("Not implemented")
        
    async def set_current_collection(self, collection_name: str):
        self.db_client.set_collection(collection_name)
        return await self.db_client.init_collection()
    
    @abstractmethod
    async def add_doc(self, path: str) -> None:
        pass

    @abstractmethod
    async def delete_doc(self, filename: str, conditions) -> None:
        pass

    async def dump_from_folder(self, dir: str, user_chat_id: str) -> None:
        """
        Add documents from a directory to the database.

        Args:
            dir (str): The directory path containing the documents.
            unstructed_client (MyUnstructuredClient): The client for unstructured data.
            embedding_client (EmbeddingClient): The client for embedding data.
            db_client (MyQdrantClient): The client for the database.

        Returns:
            None: This function does not return anything.

        This function iterates through the files in the specified directory, extracts data from each file,
        and adds the document information to the database using the provided clients.
        """

        assert os.path.isdir(dir), "Must be a valid directory"

        for path, dirnames, filenames in os.walk(dir):
            for filename in filenames:
                await self.add_doc(os.path.join(path, filename), user_chat_id)
    
    async def list_s3_files(self, session, bucket_name, prefix):
        async with session.create_client('s3', region_name="ap-southeast-1") as s3_client:
            paginator = s3_client.get_paginator('list_objects_v2')
            operation_parameters = {'Bucket': bucket_name, 'Prefix': prefix}
            page_iterator = paginator.paginate(**operation_parameters)
            
            files = []
            async for page in page_iterator:
                if "Contents" in page:
                    for item in page["Contents"]:
                        files.append(item["Key"])
        return files
    
    async def download_file(self, session, bucket_name, key, download_dir):
        async with session.create_client('s3', region_name="ap-southeast-1") as s3_client:
            try:
                response = await s3_client.get_object(Bucket=bucket_name, Key=key)
                file_path = os.path.join(download_dir, key)
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                async with response['Body'] as stream:
                    with open(file_path, 'wb') as f:
                        chunk = await stream.read()
                        f.write(chunk)
                return file_path
            except Exception as e:
                return f"Failed to download {key}: {e}"
            
    async def dump_from_s3_folder(self, dir: str, user_chat_id: str) -> None:

        session = get_session()
        files_to_download = await self.list_s3_files(session, os.environ.get("S3_BUCKET_NAME"), dir)
        if len(files_to_download) == 0:
            raise Exception("No files found in S3 bucket")
        
        tasks = [
            self.download_file(session, os.environ.get("S3_BUCKET_NAME"), key, os.environ.get("S3_DOWNLOAD_DIR"))
            for key in files_to_download
        ]
        results = await asyncio.gather(*tasks)
        failed = [result for result in results if "Failed to download" in result]
        results = [result for result in results if "Failed to download" not in result]

        if not isinstance(self.embedding_client, S3EmbeddingClient):
            for filename in results: #TODO: prevent processing of duplicate files
                await self.add_doc(filename, user_chat_id)
            
            ### When not using azure free tier, can use this instead of the above, maybe set max concurrency to <15
            # tasks = [self.add_doc(filename) for filename in results] #TODO: prevent processing of duplicate files
            # _ = await asyncio.gather(*tasks)
        else:
            # has to be done in batches so need a new method
            await self.add_doc_from_s3(results, user_chat_id)

        return failed

    async def dump_from_s3(self, file_list: List[str], user_chat_id: str, is_chunked: bool) -> None:

        session = get_session()
        tasks = [
            self.download_file(session, os.environ.get("S3_BUCKET_NAME"), key, os.environ.get("S3_DOWNLOAD_DIR"))
            for key in file_list
        ]
        print("starting to download")
        results = await asyncio.gather(*tasks)
        download_failed = [result for result in results if "Failed to download" in result]
        download_success = [result for result in results if "Failed to download" not in result]
        
        print("starting to process")
        success, char_counts = [], []

        for file in download_success:
            path, char_count = await self.add_doc_from_s3(file, user_chat_id, is_chunked)
            char_counts.append(char_count)
            success.append(path)

        return success, char_counts