from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import AsyncIterator, Dict, Generic, List, Optional, Tuple, TypeVar

import yaml
from langchain_community.vectorstores.faiss import FAISS
from langchain_core.embeddings import Embeddings
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_core.language_models import BaseChatModel, BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import BasePromptTemplate
from langchain_core.pydantic_v1 import BaseModel
from langchain_core.runnables import Runnable
from langchain_core.tracers.context import collect_runs
from langchain_core.vectorstores import VectorStore

from docugami_langchain.chains.params import ChainParameters
from docugami_langchain.config import (
    DEFAULT_EXAMPLES_PER_PROMPT,
    MAX_PARAMS_CUTOFF_LENGTH_CHARS,
)
from docugami_langchain.output_parsers import KeyfindingOutputParser
from docugami_langchain.prompts import (
    chat_prompt_template,
    generic_string_prompt_template,
)

T = TypeVar("T")

CONFIG_KEY: str = "config"
RUN_NAME_KEY: str = "run_name"


@dataclass
class TracedChainResponse(Generic[T]):
    value: T
    run_id: str = ""


class BaseDocugamiChain(BaseModel, Generic[T], ABC):
    """
    Base class with common functionality for various chains.
    """

    llm: BaseLanguageModel
    embeddings: Embeddings
    examples_vectorstore_cls: type[VectorStore] = FAISS

    input_params_max_length_cutoff: int = MAX_PARAMS_CUTOFF_LENGTH_CHARS
    few_shot_params_max_length_cutoff: int = MAX_PARAMS_CUTOFF_LENGTH_CHARS
    _examples: list[dict] = []
    _example_selector: Optional[SemanticSimilarityExampleSelector] = None

    class Config:
        arbitrary_types_allowed = True
        underscore_attrs_are_private = True

    def vector_collection_name(self) -> str:
        """
        Unique vector collection name for each class and embedding.
        """
        embedding_model_name = getattr(
            self.embeddings,
            "model_name",
            getattr(self.embeddings, "model", getattr(self.embeddings, "name", None)),
        )
        if not embedding_model_name:
            raise Exception(f"Could not determine model name for {self.embeddings}")

        raw_name = str(self.__class__.__name__)
        raw_name = f"{raw_name}-{embedding_model_name}"
        return "".join([char for char in raw_name if char.isalnum()])

    def load_examples(
        self,
        examples_yaml: Path,
        num_examples: int = DEFAULT_EXAMPLES_PER_PROMPT,
    ) -> None:
        """
        Optional: loads examples from the given examples file (YAML) and initializes an
        internal example selector that select appropriate examples for each prompt. Note
        that each chain requires its own particular format for examples based on the
        keys required in its prompt.

        The provided embeddings instance is used to embed the examples for similarity.
        See https://langchain.readthedocs.io/en/latest/modules/indexes/examples/embeddings.html.
        """

        if not self.embeddings:
            raise Exception("Embedding model required to use few shot examples")

        with open(examples_yaml, "r", encoding="utf-8") as in_f:
            self._examples = yaml.safe_load(in_f)

            for ex in self._examples:
                # truncate example length to avoid overflowing context too much
                keys = ex.keys()
                for k in keys:
                    ex[k] = ex[k][: self.few_shot_params_max_length_cutoff].strip()

            if self._examples and num_examples:
                try:
                    self._example_selector = (
                        SemanticSimilarityExampleSelector.from_examples(
                            examples=self._examples,
                            embeddings=self.embeddings,
                            vectorstore_cls=self.examples_vectorstore_cls,
                            k=num_examples,
                        )
                    )
                except Exception as exc:
                    details = f"Exception while loading samples from YAML {examples_yaml}. Details: {exc}"
                    raise Exception(details)

    def prompt(
        self,
        chain_params: ChainParameters,
        num_examples: int = DEFAULT_EXAMPLES_PER_PROMPT,
    ) -> BasePromptTemplate:
        if isinstance(self.llm, BaseChatModel):
            # For chat model instances, use chat prompts with
            # specially crafted system and few shot messages.
            return chat_prompt_template(
                chain_params=chain_params,
                example_selector=self._example_selector,
                num_examples=min(num_examples, len(self._examples)),
            )
        else:
            # For non-chat model instances, we need a string prompt
            return generic_string_prompt_template(
                chain_params=chain_params,
                example_selector=self._example_selector,
                num_examples=min(num_examples, len(self._examples)),
            )

    def runnable(self) -> Runnable:
        """
        Runnable for this chain, built dynamically from the params
        """

        # Build up prompt for this use case, possibly customizing for this model
        params = self.chain_params()
        prompt_template = self.prompt(params)

        # Generate answer from the LLM
        full_runnable = prompt_template | self.llm.bind(stop=params.stop_sequences)
        if isinstance(self.llm, BaseChatModel):
            # For chat models, we need to make sure the output is a string
            full_runnable = full_runnable | StrOutputParser()

        if params.key_finding_output_parse:
            # Increase accuracy for models that require very specific output, by
            # looking for the output key however adding such an output parser disables
            # streaming, so use carefully
            full_runnable = full_runnable | KeyfindingOutputParser(
                output_key=params.output.key
            )

        if params.additional_runnables:
            for runnable in params.additional_runnables:
                full_runnable = full_runnable | runnable

        return full_runnable

    def _prepare_run_args(self, kwargs_dict: Dict) -> Tuple[Dict, Dict]:
        # In langsmith, default the run to be named according the the chain class
        config = {RUN_NAME_KEY: self.__class__.__name__}
        if kwargs_dict and CONFIG_KEY in kwargs_dict:
            if kwargs_dict[CONFIG_KEY]:
                # Use additional caller specified config, e.g. in case of chains
                # nested inside lambdas
                config.update(**kwargs_dict[CONFIG_KEY])

            # kwargs are used as inputs to the chain prompt, so remove the config
            # param if specified
            del kwargs_dict[CONFIG_KEY]

        for key in kwargs_dict:
            # for string args, cap at max to avoid chance of prompt overflow
            if isinstance(kwargs_dict[key], str):
                kwargs_dict[key] = kwargs_dict[key][
                    : self.input_params_max_length_cutoff
                ]

        return config, kwargs_dict

    @abstractmethod
    def run(self, **kwargs) -> T:  # type: ignore
        config, kwargs_dict = self._prepare_run_args(kwargs)

        return self.runnable().invoke(input=kwargs_dict, config=config)  # type: ignore

    def traced_run(self, **kwargs) -> TracedChainResponse[T]:  # type: ignore
        with collect_runs() as cb:
            chain_output: T = self.run(**kwargs)
            run_id = str(cb.traced_runs[0].id)
            return TracedChainResponse[T](run_id=run_id, value=chain_output)

    @abstractmethod
    async def run_stream(self, **kwargs) -> AsyncIterator[TracedChainResponse[T]]:  # type: ignore
        config, kwargs_dict = self._prepare_run_args(kwargs)

        with collect_runs() as cb:
            incremental_answer = None
            async for chunk in self.runnable().astream(
                input=kwargs_dict,
                config=config,  # type: ignore
            ):
                if not incremental_answer:
                    incremental_answer = chunk
                else:
                    incremental_answer += chunk

                yield TracedChainResponse[T](value=incremental_answer)

            # yield the final result with the run_id
            if cb.traced_runs:
                run_id = str(cb.traced_runs[0].id)
                yield TracedChainResponse[T](
                    run_id=run_id,
                    value=incremental_answer,  # type: ignore
                )

    @abstractmethod
    def run_batch(self, **kwargs) -> list[T]:  # type: ignore
        config, kwargs_dict = self._prepare_run_args(kwargs)

        inputs = kwargs_dict.get("inputs")
        if not inputs:
            raise Exception("Please specify a batch for inference")

        if not isinstance(inputs, List):
            raise Exception("Input for batch processing must be a List")

        for input_dict in inputs:
            for key in input_dict:
                # For string args, cap at max to avoid chance of prompt overflow
                if isinstance(input_dict[key], str):
                    input_dict[key] = input_dict[key][
                        : self.input_params_max_length_cutoff
                    ]

        return self.runnable().batch(inputs=inputs, config=config)  # type: ignore

    @abstractmethod
    def chain_params(self) -> ChainParameters: ...
