"""Classes to simplify running queries against Claude3 using the Anthropic
API. These are for tabulr style queries against claude, i.e.
'classify each row in this input data' type inputs.
"""
from pyaddons import utils
from pyaddons.flat_files import header
from tqdm import tqdm
from hashlib import md5
import re
import anthropic
import os
import csv
import shelve
# import pprint as pp


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class BaseClaude(object):
    """Use the Claude3 API to annotate data and provide output in some tabular
    format.

    Parameters
    ----------
    api_key : `str`
        Your Anthropic API key to use for Claude requests. If not provided
        will default to ``ANTHROPIC_API_KEY``
    model : `str`
        The Claude model you want to use, this has only ever been used with
        Opus and I would not recommend using and other models.
    max_tokens : `int`, optional, default: `1024`
        The max tokens to return from an API call.
    temperature : `float`, optional, default: `0`
        The temperature setting for verbosity of Claude output, this should be
        0 as this is an extraction task.
    """
    PROMPT = None
    """The system prompt used to setup Claude for indication parsing (`str`)
    """
    HEADER = []
    """The expected header to identify when indication result data table is
    identified in the Claude results (`list` of `str`).
    """
    INTERNAL_DELIMITER = "|"
    """An internal delimiter for the data table claude returns (`str`)
    """
    TOKEN_FUDGE = 0.25
    """This is an adjustment for the the number of tokens that will be sent
    to Claude based on what we think we will get back and the max tokens to
    return. So a if max tokens is 1000 and the fudge is 0.25, then we will
    send 250 tokens on the assumption that the claude output from those tokens
    will be ~1000 (`float`).
    """
    MESSAGE_KEYWORD = "rows"
    """A key word that is used for progress messages (`str`)
    """
    API_KEY_ENV = "ANTHROPIC_API_KEY"
    """The default environment variable name for the Claude API key (`str`)
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, api_key, model, max_tokens=1024,
                 temperature=0):
        self.api_key = api_key or os.environ.get(self.API_KEY_ENV)
        self.model = model
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.client = None
        self._is_open = False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __enter__(self):
        """Enter the context manager.
        """
        return self.open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __exit__(self, *args):
        """Exit the context manager.
        """
        self.close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Open the input/output

        Returns
        -------
        self : `pyaddons.llm.claude.BaseClaude`
            For chaining.
        """
        # Initialise the client
        self.client = anthropic.Anthropic(
            # defaults to os.environ.get("ANTHROPIC_API_KEY")
            api_key=self.api_key,
        )
        self._is_open = True
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close the input/output
        """
        self._is_open = False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_input_count(self):
        """Get the number of rows in the input data.

        This is optional, and if overridden then will provide TQDM with a
        count of the input for progress monitoring.

        Returns
        -------
        input_rows : `NoneType`
            The input rows. This default implementation returnes nothing.
        """
        return None

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def yield_input(self):
        """Yield input rows to run against claude. This should be overridden
        in sub-classes

        Yields
        ------
        input_rows : `any`
            The input rows to Claude.

        Raises
        ------
        NotImplementedError
            The base class needs to be overridden.
        """
        raise NotImplementedError("Need to override this method")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def run(self, verbose=False, batch_size=100):
        """Run an annotation task. This controls input processing, batching
        and output, in short - the whole thing.

        Parameters
        ----------
        verbose : `bool`, optional, default: `False`
            Report progress.
        batch_size : `int`, optional, default: `100`
            The number of input rows to pass to the annotation method. This
            default seems fine for 4096. However, the annotation method will
            perform it's own batching as well.
        """
        tqdm_kwargs = dict(
            desc=f"[info] processing {self.MESSAGE_KEYWORD}...",
            disable=not verbose,
            unit=f"{self.MESSAGE_KEYWORD}(s)",
            leave=False,
            total=self.get_input_count()
        )

        # Will hold the input data for a batch, a batch is a single
        # submission to claude
        batch = []

        # Loop through input data
        for i in tqdm(self.yield_input(), **tqdm_kwargs):
            # Store the drug effect ID that will link results with the drug
            # effects and the actual input data
            batch.append((len(batch) + 1, i, self.get_batch_entry(i)))

            # We have reached the batch size so issue the query, the results
            # should be the same length as the input data. Where results a
            # missing an empty list should be present
            if len(batch) >= batch_size:
                # Query, providing the drug effect strings only
                results = self.annotate([k for i, j, k in batch])

                # Any missing entries will be passed back to the next batch
                # next_batch = self.find_missing(batch, results)

                # Output the results
                self.output(batch, results)

                # Go again
                batch = []

        # Process any remaining batches
        if len(batch) > 0:
            results = self.annotate([k for i, j, k in batch])

            # Any missing entries will be passed back to the next batch
            # next_batch = self.find_missing(batch, results)

            # Output the results
            self.output(batch, results)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_batch_entry(self, row):
        """Process an input row for passing to claude, this will extract the
        required data from the input row. Needs to be overridden in sub-class.

        Parameters
        ----------
        row : `any`
            An input row.

        Returns
        -------
        claude_data : `str`
            Data for claude.
        """
        raise NotImplementedError("You must overide this method")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def output(self, input_data, results):
        """Output data returned by Claude.

        This relies on the result indexes being the same as the input data
        indexes. i.e. result 0 is the result for input 0. Missing results
        should be empty lists.

        Parameters
        ----------
        input_data : `list` of `tuple`
            The input data, each input tuple should have the structure:

            0. The batch input index ID (`int`)
            1. The input string to claude.
            2. The results from claude

        results : `list` of `list`
            The results returned by claude, each sublist is a returned row
            that will match with an input row.
        """
        raise NotImplementedError("You must overide this method")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def annotate(self, strings):
        """Annotate a batch of strings.

        Parameters
        ----------
        strings : `list` of `str`
            Strings to process using claude.

        Returns
        -------
        results `list` of `list` of `list`
            Each element of the results list will correspond to an element of
            the input file. Where something has failed then an empty list will
            be present at that element. The reason for the sub-lists is to
            account for the possibility that the user has asked claude to
            place the output of an input row over several output rows with
            the same input row ID. This created a 1:many relationship between
            input and output.
        """
        # We attempt to batch up the strings into batches that are roughly
        # sized so that results will not be truncated due to hitting the max
        # token limit. This is not an exact science and some might still get
        # truncated. However, truncated results will be zeroed should should
        # get filled in on subsequent runs.
        batches = self.get_batches(strings)

        # The offset to each batch created, so batches can be assigned to the
        # correct results index
        offset = 0

        # Generate a results list that will be filled with data from Claude.
        # This is pre-created as the combination of batch data/missing results
        # may mean there will be gaps and we can't rely on just appending.
        all_results = [[] for i in strings]

        # Loop through all the batches
        for b in batches:
            # Issue the query
            try:
                message = self.client.messages.create(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    system=self.PROMPT,
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": "{0}".format(
                                        self.prep_for_message(b)
                                    )
                                }
                            ]
                        }
                    ]
                )
            except AttributeError as e:
                if self.client is None:
                    msg = "No client, are you sure you called open()?"
                    raise AttributeError(msg) from e
                raise

            # We only have a single message so only expect a single response
            # that we have to parse
            results = message.content[0].text

            # Get the reason that claude stopped generating content. Hopefully
            # this is because it reached a natural stop and not reached max
            # token count. We use this later on.
            stop_reason = message.stop_reason

            # Initial processing of the results, split the long text string
            # into result rows
            results = self.process_return_text(results)
            try:
                # Find the start of our results table
                idx = self.get_data_start_index(results)
            except IndexError:
                # Nothing found
                continue
            try:
                # Seed the last_index with the first row (but do not process)
                # (list copy).
                row = self.process_row(list(results[idx]))
                last_idx = row[0]

                # This accounts for the index of the starting row that claude
                # has returned. In most cases it will be 1-based so this will
                # be substracted from the index
                idx_offset = last_idx
            except (IndexError, ValueError, KeyError):
                # There is nothing in the results so just move onto the next
                # batch
                continue

            # Will hold all the annotations relating to a single input row
            # note there maybe > 1 annotation row per input row if claude has
            # been asked to separate over several rows. Therefore, they are
            # in nested list
            cur_row = []
            # Loop through the results table
            for row in results[idx:]:
                try:
                    # Format the indication row
                    self.process_row(row)

                    # Have we reached a new input row?
                    if row[0] != last_idx:
                        # Yes, so store the results and update the last index
                        all_results[last_idx+offset-idx_offset] = cur_row
                        cur_row = []
                        last_idx = row[0]
                    cur_row.append(row)
                except (IndexError, ValueError, KeyError):
                    # Truncated data, probably due to reaching max tokens
                    stop_reason = "error"
                    break

            try:
                # Get any data rows that have still to be processed
                if len(cur_row) > 0:
                    all_results[last_idx + offset - idx_offset] = cur_row
            except (IndexError, ValueError, KeyError):
                # Truncated data, probably due to reaching max tokens
                stop_reason = "error"

            # Claude stopped due to a non-natural stopping point, so we bin
            # the last entry as it is likely incomplete
            if stop_reason != "end_turn":
                try:
                    all_results[last_idx+offset-idx_offset] = []
                except IndexError:
                    raise

            # Adjust the offset for the next batch
            offset += len(b)
        return all_results

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def process_return_text(cls, text):
        """Process the text returned by Claude. This will break on new lines
        and split the data on some internal delimiter.

        Parameters
        ----------
        text : `str`
            The text returned by Claude.

        Returns
        -------
        data : `list` of `list`
            A provisional data table of results from Claude.
        """
        return [
            i.split(cls.INTERNAL_DELIMITER) for i in text.split("\n")
        ]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def get_data_start_index(cls, data):
        """Claude may insert text before the data table starts. This will
        identify the starting index for the actual data we are interested in.

        Parameters
        ----------
        data : `list` of `list`
            The split data returned by Claude.

        Returns
        -------
        start_index : `int`
            The start index of the data (the first row after any header).
        """
        # Find the start of our results table
        idx = 0
        while data[idx] != cls.HEADER:
            idx += 1
        idx += 1
        return idx

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def check_claude_row(self, row):
        """Check a row returned by claude to see if it is the expected length.
        The expected length is the same length as the header expected back
        from claude.

        Parameters
        ----------
        row : `list` of `str`
            A row returned by claude.

        Raises
        ------
        IndexError
            If the returned row is not the same length as the expetced header.
        """
        # Flag incomplete outputs
        if len(row) != len(self.HEADER):
            raise IndexError("Incomplete output")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def process_row(self, row):
        """Process a row of Claude output.

        Parameters
        ----------
        row : `list` of `str`
            A row to process.

        Returns
        -------
        row : `list`
            A processed row. It is expected that the first element of a
            processed claude row is an index value that can be used to link
            the rows returned by Claude to rows given to claude.
        """
        self.check_claude_row(row)
        row[0] = int(row[0])
        return row

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_batches(self, strings):
        """Batch up the strings into submission batches, using the max tokens
        as a guide. We will submit ``token_fudge`` the number of input tokens.

        Parameters
        ----------
        strings : `list` of `str`
            The string data we want Claude to process.

        Returns
        -------
        batches : `list` of `list` of `str`
            The batches we will run against claude.
        """
        batches = []
        curr_batch = []
        target = self.max_tokens * self.TOKEN_FUDGE
        cur_tokens = 0
        for i in strings:
            if i.strip() == "":
                # Skip any completely empty strings
                continue
            i = re.sub('[\r\n\t]', ' ', i)
            tokens = i.split(" ")
            if cur_tokens + len(tokens) < target:
                curr_batch.append(i)
                cur_tokens += len(tokens)
            else:
                batches.append(curr_batch)
                curr_batch = [i]
                cur_tokens = len(tokens)
        if len(curr_batch) > 0:
            batches.append(curr_batch)
        return batches

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def prep_for_message(cls, input_rows):
        """Prepare the input rows for attaching to the Claude prompt.

        The default implementation is just to join on new line, but it may be
        better to enclose in XML tags (see XmlInputMixin)

        Parameters
        ----------
        input_rows : `list` of `str`
            The input rows.

        Returns
        -------
        message_str : `str`
            The message string for the Claude prompt.
        """
        return "\n".join(input_rows)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class ClaudeFile(BaseClaude):
    """Use the Claude3 API to annotate data that is in an input file and
    output to an output file. This handles output caching/interuption pick up

    Parameters
    ----------
    infile : `str`
        The path to the input file containing the prescription strings.
    outfile : `str`
        The path to the output file of finsihed classifications. Note output
        is written via temp.
    api_key : `str`
        Your Anthropic API key to use for Claude requests.
    model : `str`
        The Claude model you want to use, this has only ever been used with
        Opus and I would not recommend using and other models.
    max_tokens : `int`, optional, default: `1024`
        The max tokens to return from an API call.
    temperature : `float`, optional, default: `0`
        The temperature setting for verbosity of Claude output, this should be
        0 as this is an extraction task.
    delimiter : `str`, optional, defafult: `\t`
        The delimiter of the input/output files.
    skiplines : `int`, optional, default: `0`
        Skip this many fixed lines at the start of the file before processing.
    comment : `str`, optional, default: `NoneType`
        Ignore any lines starting with this comment string.
    tmpdir : `str`, optional, default: `NoneType`
        A directory to write the temp file. If not set this will default to
        the system temp location.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, infile, outfile,  *args, delimiter="\t", skiplines=0,
                 comment=None, tmpdir=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.infile = infile
        self.outfile = outfile
        self.skiplines = skiplines
        self.comment = comment
        self.tmpdir = tmpdir
        self.delimiter = delimiter

        self._input_header = None
        self.input_open_method = None
        self.out_open_method = None
        self._shelve = None
        self._input_rows = 0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Initialise the shelve DB for tmp row store/
        """
        # Initialise the input file header
        h = self.input_header

        # Initialise the shevle database and call super open
        self._open_shelve()
        return super().open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _open_shelve(self):
        """Open the temp shelve database for pickup.
        """
        # Get the location of the temp directory, if it is None, this will
        # reveal the system temp location
        test_tmp = utils.get_tmp_file(dir=self.tmpdir)
        tmpdir = os.path.dirname(test_tmp)
        os.unlink(test_tmp)
        prompt_md5 = md5(self.PROMPT.encode()).hexdigest()
        self.tmpfile = os.path.join(
            tmpdir,
            f"{utils.get_file_md5(self.infile)}-{prompt_md5}.claude.shelve"
        )

        self._shelve = shelve.open(self.tmpfile)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close the sheve database and write output of all input rows have
        been processed. Sheve DB is deleted on success but kept on failure.
        """
        # Have all the input rows been entered into the shelve database
        if len(self._shelve) == self._input_rows:
            self.out_open_method = utils.get_open_method(str(self.outfile))
            with self.out_open_method(self.outfile, 'wt') as outfile:
                writer = csv.DictWriter(
                    outfile, self.get_output_header(),
                    delimiter=self.delimiter
                )
                writer.writeheader()
                for i in sorted(self._shelve.keys(), key=lambda x: int(x)):
                    for x in self._shelve[i]:
                        writer.writerow(x)
            self._shelve.close()

            # Delete the sheve data
            os.unlink(f"{self.tmpfile}.bak")
            os.unlink(f"{self.tmpfile}.dat")
            os.unlink(f"{self.tmpfile}.dir")
        else:
            self._shelve.close()
        super().close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def input_header(self):
        """Get the input file header, this will lazy load the header if not
        set (`list` of `str`)
        """
        if self._input_header is not None:
            return self._input_header

        om = self.input_open_method or utils.get_open_method(self.infile)
        with om(self.infile) as infile:
            header.move_to_header(infile, self.skiplines, self.comment)
            self._input_header = next(
                csv.reader([next(infile)], delimiter=self.delimiter)
            )
        self.input_open_method = om
        return self._input_header

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_input_count(self):
        """Get the count of input rows that need processing. This will be
        adjusted to factor in rows that have already been processed.

        Returns
        -------
        input_count : `int`
            The number of rows to be processed.
        """
        if self._input_rows > 0:
            return self._input_rows - len(self._shelve)
        om = self.input_open_method or utils.get_open_method(self.infile)
        with om(self.infile) as intmp:
            header.move_to_header(intmp, self.skiplines, self.comment)
            h = next(intmp)
            for idx, row in enumerate(intmp, 1):
                pass
        self._input_rows = idx
        return self._input_rows - len(self._shelve)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_output_header(self):
        """Get the output file header, this is the input file header and the
        expected claude header.

        Returns
        -------
        output_header : `list` of `str`
        """
        return self.input_header + self.HEADER

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def yield_input(self):
        """Yield input file rows for processing.

        This will skip input rows that are present in the shelve database.
        So only unprocessed rows are yielded.

        Yields
        ------
        row_id : `str`
            This is the input row number cast to a string to be compatible
            with shelve
        input_row : `dict`
            An unprocessed input row.
        """
        try:
            with self.input_open_method(self.infile) as infile:
                header.move_to_header(infile, self.skiplines, self.comment)
                reader = csv.DictReader(infile, delimiter=self.delimiter)

                for idx, row in enumerate(reader, 1):
                    if str(idx) not in self._shelve:
                        yield str(idx), row
        except TypeError as e:
            if self.input_open_method is None:
                raise TypeError("Did you open the object?") from e
            raise

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def output(self, input_data, results):
        """Process the output data returned by Claude.

        This relies on the result indexes being the same as the input data
        indexes. i.e. result 0 is the result for input 0. Missing results
        should be empty lists.

        Parameters
        ----------
        input_data : `list` of `tuple`
            The input data, each input tuple should have the structure:

            0. The batch input index ID (`int`)
            1. A tuple of the input row_id (`str`) and the input row (`dict`)
            2. The input string to claude.

        results : `list` of `list` of `list`
            The results returned by claude, each sublist is a returned row
            that will match with an input row.

        Notes
        -----
        This will store result rows that have data (length > 0) into the sheve
        database.
        """
        for i, r in zip(input_data, results):
            _, inrow, annot_str = i
            shelve_idx, inrow = inrow
            if len(r) == 0:
                continue

            store_list = []
            for x in r:
                store_list.append(
                    {
                        **inrow,
                        **dict(
                            [(k, v) for k, v in zip(self.HEADER, x)]
                        )
                    }
                )
            self._shelve[shelve_idx] = store_list


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class XmlInputMixin(object):
    """A mixin that wraps <input></input> XML tags around each input row to
     claude.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def prep_for_message(cls, input_rows):
        """Prepare the input rows for attaching to the Claude prompt.

        Parameters
        ----------
        input_rows : `list` of `str`
            The input rows.

        Returns
        -------
        message_str : `str`
            The message string for the Claude prompt.
        """
        return "\n".join(
            [f"<input>\n{i}\n</input>\n" for i in input_rows]
        )
