from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, List

import pandas as pd
import streamlit as st
from langchain.chains import LLMChain
from langchain.prompts.few_shot import FewShotPromptTemplate

from doccano_mini.components import (
    display_download_button,
    openai_model_form,
    task_instruction_editor,
    usage,
)
from doccano_mini.utils import escape_markdown


class BasePage(ABC):
    example_path: str = ""

    def __init__(self, title: str) -> None:
        self.title = title

    @property
    def columns(self) -> List[str]:
        return []

    def load_examples(self, filename: str) -> pd.DataFrame:
        filepath = Path(__file__).parent.resolve().joinpath("examples", filename)
        return pd.read_json(filepath)

    def make_examples(self, columns: List[str]) -> List[Dict]:
        df = self.load_examples(self.example_path)
        edited_df = st.experimental_data_editor(df, num_rows="dynamic", width=1000)
        examples = edited_df.to_dict(orient="records")
        return examples

    @abstractmethod
    def make_prompt(self, examples: List[Dict]) -> FewShotPromptTemplate:
        raise NotImplementedError()

    @abstractmethod
    def prepare_inputs(self, columns: List[str]) -> Dict:
        raise NotImplementedError()

    def annotate(self, examples: List[Dict]) -> List[Dict]:
        return examples

    def render(self) -> None:
        st.title(self.title)
        st.header("Annotate your data")
        columns = self.columns
        examples = self.make_examples(columns)
        examples = self.annotate(examples)

        prompt = self.make_prompt(examples)
        prompt = task_instruction_editor(prompt)

        st.header("Test")
        col1, col2 = st.columns([3, 1])

        with col1:
            inputs = self.prepare_inputs(columns)

        with col2:
            llm = openai_model_form()

        with st.expander("See your prompt"):
            st.markdown(f"```\n{prompt.format(**inputs)}\n```")

        if llm is None:
            st.error("Enter your API key.")

        if st.button("Predict", disabled=llm is None):
            chain = LLMChain(llm=llm, prompt=prompt)  # type:ignore
            response = chain.run(**inputs)
            st.markdown(escape_markdown(response).replace("\n", "  \n"))

            chain.save("config.yaml")
            display_download_button()
        usage()

langchain.chains.LLMChain
[((763, 785), 'pandas.read_json', 'pd.read_json', (['filepath'], {}), '(filepath)\n', (775, 785), True, 'import pandas as pd\n'), ((921, 984), 'streamlit.experimental_data_editor', 'st.experimental_data_editor', (['df'], {'num_rows': '"""dynamic"""', 'width': '(1000)'}), "(df, num_rows='dynamic', width=1000)\n", (948, 984), True, 'import streamlit as st\n'), ((1434, 1454), 'streamlit.title', 'st.title', (['self.title'], {}), '(self.title)\n', (1442, 1454), True, 'import streamlit as st\n'), ((1463, 1494), 'streamlit.header', 'st.header', (['"""Annotate your data"""'], {}), "('Annotate your data')\n", (1472, 1494), True, 'import streamlit as st\n'), ((1678, 1709), 'doccano_mini.components.task_instruction_editor', 'task_instruction_editor', (['prompt'], {}), '(prompt)\n', (1701, 1709), False, 'from doccano_mini.components import display_download_button, openai_model_form, task_instruction_editor, usage\n'), ((1719, 1736), 'streamlit.header', 'st.header', (['"""Test"""'], {}), "('Test')\n", (1728, 1736), True, 'import streamlit as st\n'), ((1758, 1776), 'streamlit.columns', 'st.columns', (['[3, 1]'], {}), '([3, 1])\n', (1768, 1776), True, 'import streamlit as st\n'), ((2096, 2138), 'streamlit.button', 'st.button', (['"""Predict"""'], {'disabled': '(llm is None)'}), "('Predict', disabled=llm is None)\n", (2105, 2138), True, 'import streamlit as st\n'), ((2409, 2416), 'doccano_mini.components.usage', 'usage', ([], {}), '()\n', (2414, 2416), False, 'from doccano_mini.components import display_download_button, openai_model_form, task_instruction_editor, usage\n'), ((1885, 1904), 'doccano_mini.components.openai_model_form', 'openai_model_form', ([], {}), '()\n', (1902, 1904), False, 'from doccano_mini.components import display_download_button, openai_model_form, task_instruction_editor, usage\n'), ((1919, 1949), 'streamlit.expander', 'st.expander', (['"""See your prompt"""'], {}), "('See your prompt')\n", (1930, 1949), True, 'import streamlit as st\n'), ((2052, 2083), 'streamlit.error', 'st.error', (['"""Enter your API key."""'], {}), "('Enter your API key.')\n", (2060, 2083), True, 'import streamlit as st\n'), ((2160, 2192), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (2168, 2192), False, 'from langchain.chains import LLMChain\n'), ((2375, 2400), 'doccano_mini.components.display_download_button', 'display_download_button', ([], {}), '()\n', (2398, 2400), False, 'from doccano_mini.components import display_download_button, openai_model_form, task_instruction_editor, usage\n'), ((2275, 2300), 'doccano_mini.utils.escape_markdown', 'escape_markdown', (['response'], {}), '(response)\n', (2290, 2300), False, 'from doccano_mini.utils import escape_markdown\n'), ((685, 699), 'pathlib.Path', 'Path', (['__file__'], {}), '(__file__)\n', (689, 699), False, 'from pathlib import Path\n')]