# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2023
# SPDX-License-Identifier: Apache-2.0

from typing import Any, Dict, List, Optional

from langchain.agents import tool
from langchain.chains.base import Chain
from langchain.chains import LLMChain
from langchain import PromptTemplate
from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)

import chainlit as cl
from chainlit.context import context
from chainlit import run_sync

from tabulate import tabulate

from ..llm import get_bedrock_text, get_processed_prompt_template
from .graph import GraphChain


def get_tool_metadata():
    return {
        "name": "3dview",
        "description": "Useful to teleport in 3D viewer to the equipment the user is interested in. \
            Input to this tool should be the entityId of the equipment. \
            Output is a string to confirm whether the view is found or not.",
    }

@tool
def run(input: str) -> str:
    """Identify the location of the object user is asking about."""

    point_camera_to_entity(input)
    return 'Found it!'

def point_camera_to_entity(entityId):
    run_sync(context.session.emit('view', entityId))

ENTITY_EXTRACTION_PROMPT = """
Your job is to identify the entity user is asking about based on the user question.

Use the following format:

Question: the input question from the user
Entity: the phrase about the entity in the original question

Only output the entity phrase, do not repeat the question.

Here are some examples:

Question: teleport me to the cookie line in alarm state
Entity: the cookie line in alarm state

Question: show me the freezer tunnel
Entity: the freezer tunnel

Question: show me the conveyer belt
Entity: the conveyer belt

Now begin!

Question: {question}
Entity:
"""

class EntityExtractorChain(Chain):
    """Chain to find the entity in the question."""
    
    llm_chain: LLMChain

    @property
    def input_keys(self) -> List[str]:
        return ['question']

    @property
    def output_keys(self) -> List[str]:
        return ['entity']
    
    @classmethod
    def create(cls, **kwargs):
        llm = get_bedrock_text()
        prompt = PromptTemplate(
            template=get_processed_prompt_template(ENTITY_EXTRACTION_PROMPT),
            input_variables=["question"],
        )
        llm_chain = LLMChain(
            llm=llm,
            prompt=prompt,
            **kwargs)
        return cls(llm_chain=llm_chain)

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        callbacks = _run_manager.get_child()
        output = self.llm_chain.run(callbacks=callbacks, **inputs)
        return {
            'entity': output
        }

    async def _acall(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        callbacks = _run_manager.get_child()
        output = await self.llm_chain.arun(callbacks=callbacks, **inputs)
        return {
            'entity': output
        }


class ViewChain(Chain):
    """Chain that manipulates 3D viewer."""
    
    entity_extractor: EntityExtractorChain
    entity_lookup: GraphChain

    @property
    def input_keys(self) -> List[str]:
        return ['question']

    @property
    def output_keys(self) -> List[str]:
        return ['text', 'selected_entity']
    
    @classmethod
    def create(cls, **kwargs):
        entity_extractor = EntityExtractorChain.create(**kwargs)
        entity_lookup = GraphChain.create(**kwargs)
        return cls(entity_extractor=entity_extractor, entity_lookup=entity_lookup, **kwargs)
        
    def pick_entity(self, entities):
        if entities.shape[0] > 1:
            headers = ['No', 'Name', 'Id']
            rows = [[i + 1, row.entityName, row.entityId] for i, row in entities.items()]
            entity_table = tabulate(rows, headers=headers, tablefmt="pipe")
            
            run_sync(cl.Message(content="I've found these matching entities:\n\n" + entity_table).send())
            
            res = run_sync(cl.AskUserMessage(content="Which one do you mean?").send())
            
            if res is not None:
                # TODO: use a LLMChain to parse the user input
                idx = int(res['content']) - 1
                
                entityId = entities.iloc[idx].entityId
            else:
                entityId = None
        else:
            entityId = entities.iloc[0].entityId
            
        return entityId

    async def apick_entity(self, entities):
        if entities.shape[0] > 1:
            headers = ['No', 'Name', 'Id']
            rows = [[i + 1, row.entityName, row.entityId] for i, row in entities.items()]
            entity_table = tabulate(rows, headers=headers, tablefmt="pipe")
            
            await cl.Message(content="I've found these matching entities:\n\n" + entity_table).send()
            
            res = await cl.AskUserMessage(content="Which one do you mean?").send()
            
            if res is not None:
                # TODO: use a LLMChain to parse the user input
                idx = int(res['content']) - 1
                
                entityId = entities.iloc[idx].entityId
            else:
                entityId = None
        else:
                entityId = entities.iloc[0].entityId
            
        return entityId

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        callbacks = _run_manager.get_child()
        
        entity = self.entity_extractor.run(callbacks=callbacks, **inputs)
        df = self.entity_lookup.run(
            callbacks,
            {
                "question": "Find all entities matching the description: " + entity,
                "format_output": False
            })
        
        # TODO: handle the column detection better
        if df.shape[0] < 1 or df.columns[0] != 'e':
            return {
                'text': "I didn't find any result.",
                'selected_entity': ''
            }
            
        entities = df[df.columns[0]]
        
        entityId = self.pick_entity(entities)
        
        if entityId is None:
            return {
                'text': "I didn't find any result.",
                'selected_entity': ''
            }
            
        point_camera_to_entity(entityId)
        
        return {
            'text': f"I've pointed you to the {entityId} in the 3D Viewer.",
            'selected_entity': entityId
        }

    async def _acall(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        callbacks = _run_manager.get_child()
        
        entity = await self.entity_extractor.arun(callbacks=callbacks, **inputs)
        df = await self.entity_lookup.arun(
            **{
                "question": "Find all entities matching the description: " + entity,
                "format_output": False
            })
        
        # TODO: handle the column detection better
        if df.shape[0] < 1 or df.columns[0] != 'e':
            return {
                'text': "I didn't find any result.",
                'selected_entity': ''
            }
            
        entities = df[df.columns[0]]
        
        entityId = await self.apick_entity(entities)
        
        if entityId is None:
            return {
                'text': "I didn't find any result.",
                'selected_entity': ''
            }
            
        point_camera_to_entity(entityId)
        return {
            'text': f"I've pointed you to the {entityId} in the 3D Viewer.",
            'selected_entity': entityId
        }

langchain.chains.LLMChain langchain.callbacks.manager.CallbackManagerForChainRun.get_noop_manager
[((1178, 1216), 'chainlit.context.context.session.emit', 'context.session.emit', (['"""view"""', 'entityId'], {}), "('view', entityId)\n", (1198, 1216), False, 'from chainlit.context import context\n'), ((2370, 2412), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt, **kwargs)\n', (2378, 2412), False, 'from langchain.chains import LLMChain\n'), ((2681, 2726), 'langchain.callbacks.manager.CallbackManagerForChainRun.get_noop_manager', 'CallbackManagerForChainRun.get_noop_manager', ([], {}), '()\n', (2724, 2726), False, 'from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun\n'), ((3098, 3143), 'langchain.callbacks.manager.CallbackManagerForChainRun.get_noop_manager', 'CallbackManagerForChainRun.get_noop_manager', ([], {}), '()\n', (3141, 3143), False, 'from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun\n'), ((4150, 4198), 'tabulate.tabulate', 'tabulate', (['rows'], {'headers': 'headers', 'tablefmt': '"""pipe"""'}), "(rows, headers=headers, tablefmt='pipe')\n", (4158, 4198), False, 'from tabulate import tabulate\n'), ((5033, 5081), 'tabulate.tabulate', 'tabulate', (['rows'], {'headers': 'headers', 'tablefmt': '"""pipe"""'}), "(rows, headers=headers, tablefmt='pipe')\n", (5041, 5081), False, 'from tabulate import tabulate\n'), ((5864, 5909), 'langchain.callbacks.manager.CallbackManagerForChainRun.get_noop_manager', 'CallbackManagerForChainRun.get_noop_manager', ([], {}), '()\n', (5907, 5909), False, 'from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun\n'), ((7168, 7213), 'langchain.callbacks.manager.CallbackManagerForChainRun.get_noop_manager', 'CallbackManagerForChainRun.get_noop_manager', ([], {}), '()\n', (7211, 7213), False, 'from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun\n'), ((4233, 4309), 'chainlit.Message', 'cl.Message', ([], {'content': '("I\'ve found these matching entities:\\n\\n" + entity_table)'}), '(content="I\'ve found these matching entities:\\n\\n" + entity_table)\n', (4243, 4309), True, 'import chainlit as cl\n'), ((4358, 4409), 'chainlit.AskUserMessage', 'cl.AskUserMessage', ([], {'content': '"""Which one do you mean?"""'}), "(content='Which one do you mean?')\n", (4375, 4409), True, 'import chainlit as cl\n'), ((5113, 5189), 'chainlit.Message', 'cl.Message', ([], {'content': '("I\'ve found these matching entities:\\n\\n" + entity_table)'}), '(content="I\'ve found these matching entities:\\n\\n" + entity_table)\n', (5123, 5189), True, 'import chainlit as cl\n'), ((5234, 5285), 'chainlit.AskUserMessage', 'cl.AskUserMessage', ([], {'content': '"""Which one do you mean?"""'}), "(content='Which one do you mean?')\n", (5251, 5285), True, 'import chainlit as cl\n')]