import asyncio
import copy
import json
import logging
import uuid
from typing import Any, AsyncGenerator, Optional

from r2r.base import (
    AsyncState,
    Extraction,
    Fragment,
    FragmentType,
    KGExtraction,
    KGProvider,
    KVLoggingSingleton,
    LLMProvider,
    PipeType,
    PromptProvider,
    TextSplitter,
    extract_entities,
    extract_triples,
    generate_id_from_label,
)
from r2r.base.pipes.base_pipe import AsyncPipe

logger = logging.getLogger(__name__)


class ClientError(Exception):
    """Base class for client connection errors."""

    pass


class KGExtractionPipe(AsyncPipe):
    """
    Embeds and stores documents using a specified embedding model and database.
    """

    def __init__(
        self,
        kg_provider: KGProvider,
        llm_provider: LLMProvider,
        prompt_provider: PromptProvider,
        text_splitter: TextSplitter,
        kg_batch_size: int = 1,
        id_prefix: str = "demo",
        pipe_logger: Optional[KVLoggingSingleton] = None,
        type: PipeType = PipeType.INGESTOR,
        config: Optional[AsyncPipe.PipeConfig] = None,
        *args,
        **kwargs,
    ):
        """
        Initializes the embedding pipe with necessary components and configurations.
        """
        super().__init__(
            pipe_logger=pipe_logger,
            type=type,
            config=config
            or AsyncPipe.PipeConfig(name="default_embedding_pipe"),
        )

        self.kg_provider = kg_provider
        self.prompt_provider = prompt_provider
        self.llm_provider = llm_provider
        self.text_splitter = text_splitter
        self.kg_batch_size = kg_batch_size
        self.id_prefix = id_prefix
        self.pipe_run_info = None

    async def fragment(
        self, extraction: Extraction, run_id: uuid.UUID
    ) -> AsyncGenerator[Fragment, None]:
        """
        Splits text into manageable chunks for embedding.
        """
        if not isinstance(extraction, Extraction):
            raise ValueError(
                f"Expected an Extraction, but received {type(extraction)}."
            )
        if not isinstance(extraction.data, str):
            raise ValueError(
                f"Expected a string, but received {type(extraction.data)}."
            )
        text_chunks = [
            ele.page_content
            for ele in self.text_splitter.create_documents([extraction.data])
        ]
        for iteration, chunk in enumerate(text_chunks):
            fragment = Fragment(
                id=generate_id_from_label(f"{extraction.id}-{iteration}"),
                type=FragmentType.TEXT,
                data=chunk,
                metadata=copy.deepcopy(extraction.metadata),
                extraction_id=extraction.id,
                document_id=extraction.document_id,
            )
            yield fragment

    async def transform_fragments(
        self, fragments: list[Fragment]
    ) -> AsyncGenerator[Fragment, None]:
        """
        Transforms text chunks based on their metadata, e.g., adding prefixes.
        """
        async for fragment in fragments:
            if "chunk_prefix" in fragment.metadata:
                prefix = fragment.metadata.pop("chunk_prefix")
                fragment.data = f"{prefix}\n{fragment.data}"
            yield fragment

    async def extract_kg(
        self,
        fragment: Fragment,
        retries: int = 3,
        delay: int = 2,
    ) -> KGExtraction:
        """
        Extracts NER triples from a list of fragments with retries.
        """
        task_prompt = self.prompt_provider.get_prompt(
            self.kg_provider.config.kg_extraction_prompt,
            inputs={"input": fragment.data},
        )
        messages = self.prompt_provider._get_message_payload(
            self.prompt_provider.get_prompt("default_system"), task_prompt
        )
        for attempt in range(retries):
            try:
                response = await self.llm_provider.aget_completion(
                    messages, self.kg_provider.config.kg_extraction_config
                )

                kg_extraction = response.choices[0].message.content

                # Parsing JSON from the response
                kg_json = (
                    json.loads(
                        kg_extraction.split("```json")[1].split("```")[0]
                    )
                    if """```json""" in kg_extraction
                    else json.loads(kg_extraction)
                )
                llm_payload = kg_json.get("entities_and_triples", {})

                # Extract triples with detailed logging
                entities = extract_entities(llm_payload)
                triples = extract_triples(llm_payload, entities)

                # Create KG extraction object
                return KGExtraction(entities=entities, triples=triples)
            except (
                ClientError,
                json.JSONDecodeError,
                KeyError,
                IndexError,
            ) as e:
                logger.error(f"Error in extract_kg: {e}")
                if attempt < retries - 1:
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"Failed after retries with {e}")
                    # raise e  # Ensure the exception is raised after the final attempt

        return KGExtraction(entities={}, triples=[])

    async def _process_batch(
        self,
        fragment_batch: list[Fragment],
    ) -> list[KGExtraction]:
        """
        Embeds a batch of fragments and yields vector entries.
        """
        tasks = [
            asyncio.create_task(self.extract_kg(fragment))
            for fragment in fragment_batch
        ]
        return await asyncio.gather(*tasks)

    async def _run_logic(
        self,
        input: AsyncPipe.Input,
        state: AsyncState,
        run_id: uuid.UUID,
        *args: Any,
        **kwargs: Any,
    ) -> AsyncGenerator[KGExtraction, None]:
        """
        Executes the embedding pipe: chunking, transforming, embedding, and storing documents.
        """
        batch_tasks = []
        fragment_batch = []

        fragment_info = {}
        async for extraction in input.message:
            async for fragment in self.transform_fragments(
                self.fragment(extraction, run_id)
            ):
                if extraction.document_id in fragment_info:
                    fragment_info[extraction.document_id] += 1
                else:
                    fragment_info[extraction.document_id] = 1
                extraction.metadata["chunk_order"] = fragment_info[
                    extraction.document_id
                ]
                fragment_batch.append(fragment)
                if len(fragment_batch) >= self.kg_batch_size:
                    # Here, ensure `_process_batch` is scheduled as a coroutine, not called directly
                    batch_tasks.append(
                        self._process_batch(fragment_batch.copy())
                    )  # pass a copy if necessary
                    fragment_batch.clear()  # Clear the batch for new fragments

        logger.debug(
            f"Fragmented the input document ids into counts as shown: {fragment_info}"
        )

        if fragment_batch:  # Process any remaining fragments
            batch_tasks.append(self._process_batch(fragment_batch.copy()))

        # Process tasks as they complete
        for task in asyncio.as_completed(batch_tasks):
            batch_result = await task  # Wait for the next task to complete
            for kg_extraction in batch_result:
                yield kg_extraction