from collections import OrderedDict, defaultdict
from time import time_ns
from typing import Any, DefaultDict, Dict, List, Optional, Union
from uuid import UUID

from langchain.callbacks.base import BaseCallbackHandler
from langchain.prompts.chat import ChatPromptValue
from langchain.schema.messages import BaseMessage
from langchain.schema.output import LLMResult

from promptquality.chain_run import chain_run
from promptquality.constants.scorers import Scorers
from promptquality.set_config import set_config
from promptquality.types.chains.node_type import ChainNodeType
from promptquality.types.chains.row import ChainRow
from promptquality.types.config import Config
from promptquality.types.custom_scorer import CustomScorer
from promptquality.types.registered_scorers import RegisteredScorer


class GalileoPromptCallback(BaseCallbackHandler):
    def __init__(
        self,
        project_name: Optional[str] = None,
        run_name: Optional[str] = None,
        scorers: Optional[List[Union[Scorers, CustomScorer, RegisteredScorer, str]]] = None,
        config: Optional[Config] = None,
        **kwargs: Any,
    ) -> None:
        """
        LangChain callbackbander for logging prompts to Galileo.

        Parameters
        ----------
        project_name : str
            Name of the project to log to
        """
        # Galileo parameters.
        self.project_name = project_name
        self.run_name = run_name
        self.scorers = scorers
        self.config = config or set_config()

        # Row information.
        self.root_child_mapping: DefaultDict = defaultdict(list)
        self.parent_child_mapping: DefaultDict = defaultdict(list)
        self.child_parent_mapping: Dict[UUID, Optional[UUID]] = dict()
        self.chain_node_types: Dict[UUID, str] = dict()

        self.chain_inputs: Dict[UUID, Dict[str, Any]] = dict()
        self.serializations: Dict[UUID, Dict[str, Any]] = dict()

        self.rows: OrderedDict[UUID, ChainRow] = OrderedDict()

    def set_relationships(self, run_id: UUID, node_type: ChainNodeType, parent_run_id: Optional[UUID] = None) -> None:
        self.child_parent_mapping[run_id] = parent_run_id
        if parent_run_id:
            self.parent_child_mapping[parent_run_id].append(run_id)
        root_id = self.get_root_id(run_id)
        self.root_child_mapping[root_id].append(run_id)
        self.rows[run_id] = ChainRow(
            node_id=run_id,
            node_type=node_type,
            # -1 because the step is incremented after the callback is run.
            step=len(self.root_child_mapping[root_id]) - 1,
            chain_id=parent_run_id,
            chain_root_id=root_id,
        )

    def get_root_id(self, run_id: UUID) -> UUID:
        parent_id = self.child_parent_mapping[run_id]
        if parent_id:
            return self.get_root_id(parent_id)
        else:
            return run_id

    def mark_step_start(
        self,
        run_id: UUID,
        serialized: Dict[str, Any],
        prompt: Optional[str] = None,
        **kwargs: Dict[str, Any],
    ) -> None:
        root_id = self.get_root_id(run_id)
        self.rows[run_id].serialized = serialized
        self.rows[run_id].prompt = prompt
        self.rows[run_id].params = kwargs.get("invocation_params", dict())
        self.rows[run_id].inputs = self.chain_inputs.get(root_id, dict())

    def mark_step_end(self, run_id: UUID, response: Optional[str] = None, **kwargs: Dict[str, Any]) -> None:
        self.rows[run_id].latency = time_ns() - self.rows[run_id].creation_timestamp
        self.rows[run_id].response = response

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when LLM starts running."""
        self.set_relationships(run_id, node_type=ChainNodeType.llm, parent_run_id=parent_run_id)
        self.mark_step_start(run_id, serialized, prompts[0], **kwargs)

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when Chat Model starts running."""
        self.set_relationships(run_id, node_type=ChainNodeType.chat, parent_run_id=parent_run_id)

        chat_messages = ChatPromptValue(messages=messages[0])
        self.mark_step_start(run_id, serialized, prompt=chat_messages.to_string(), **kwargs)

    def on_llm_end(
        self, response: LLMResult, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when LLM ends running."""

        # Flatten the response and pull out the generation.
        generation = response.flatten()[0].generations[0][0]
        self.mark_step_end(run_id, response=generation.text, **kwargs)

        if response.llm_output:
            usage = response.llm_output.get("token_usage", {})
            self.rows[run_id].query_input_tokens = usage.get("prompt_tokens", 0)
            self.rows[run_id].query_output_tokens = usage.get("completion_tokens", 0)
            self.rows[run_id].query_total_tokens = usage.get("total_tokens", 0)

        if generation.generation_info:
            self.rows[run_id].finish_reason = generation.generation_info.get("finish_reason", "")

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when LLM errors."""
        del self.rows[run_id]

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when chain starts running."""
        self.set_relationships(run_id, node_type=ChainNodeType.chain, parent_run_id=parent_run_id)
        self.chain_inputs[run_id] = inputs
        self.mark_step_start(run_id, serialized, **kwargs)

    def on_chain_end(
        self, outputs: Union[str, Dict[str, Any]], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when chain ends running."""
        self.rows[run_id].finish_reason = "chain_end"
        response: Optional[str] = None
        if isinstance(outputs, str):
            response = outputs
        elif isinstance(outputs, dict):
            response = outputs.get("text")
        self.mark_step_end(run_id, response=response, **kwargs)

    def on_chain_error(
        self, error: BaseException, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
    ) -> Any:
        """Run when chain errors."""

    def finish(self) -> None:
        chain_run(list(self.rows.values()), self.project_name, self.run_name, self.scorers, config=self.config)
