import logging
import re
from functools import cache
from pathlib import Path
from typing import List, Set, Tuple, TypeVar

import torch
from PIL import Image

from utils import device, nested_apply, sorted_list

RE_PATTERN = r'^(deselect\s[A-Z](?:\s[A-Z])*(?:\sselect\s[A-Z](?:\s[A-Z])*)?|select\s[A-Z](?:\s[A-Z])*)$'  # noqa


# Name type, newtype of str. e.g. "page4-249.png"
N = TypeVar('N')

ALPHABET = 'ABCDEFGHIJ'  # we only have 10 images
LEGAL_TOKEN_IDS = [2, 315, 330, 334, 365, 382, 384, 401, 413,
                   420, 475, 5339, 634, 17960, 32002]  # A - J and <end_of_utterance> and <\s> and 'select' and 'deselect'


MINI_DECODER = {
    384: 'D',
    # 2: '</s>',
    32002: '<end_of_utterance>',
    420: 'G', 17960: 'elect',
    330: 'A', 365: 'B', 334: 'C', 5339: 'select', 401: 'F', 475: 'J',
    634: 'des', 315: 'I', 413: 'E', 382: 'H'}


class AlphabeticNameHash:

    @cache
    def __init__(self, context: List[N]) -> None:
        self._forward_map = {im: ALPHABET[i] for i, im in enumerate(context)}
        self._backward_map = {ALPHABET[i]: im for i, im in enumerate(context)}

    def hash(self, im: N) -> str:
        return self._forward_map[im]

    def unhash(self, i: str) -> N:
        return self._backward_map[i]

    def valid_hash(self, i: str) -> bool:
        return i in self._backward_map


class IdeficsAdapter:

    PAD_TOKEN_ID = 0
    LABEL_MASK_ID = 32001  # idefics2: image_token_id
    LEGAL_TOKEN_IDS = LEGAL_TOKEN_IDS
    LEGAL_TOKEN_MASK = torch.zeros(32003, requires_grad=False)\
        .index_fill_(0, torch.tensor(LEGAL_TOKEN_IDS), 1).to(device=device(), dtype=torch.bool)
    SUPPRESS_TOKEN_IDS = list(set(range(32003)) - set(LEGAL_TOKEN_IDS))

    def __init__(self, image_folder: str, processor) -> None:
        self.t_max_length = 2048
        self.image_folder = Path(image_folder)
        self.image_cache = {}
        self.processor = processor
        self.tokenizer = self.processor.tokenizer

    def get_image(self, im_name: N) -> Image.Image:
        if im_name not in self.image_cache:
            self.image_cache[im_name] = Image.open(
                self.image_folder.joinpath(im_name))
        return self.image_cache[im_name]

    def unhash(self, context: List[N], c: str):
        return AlphabeticNameHash(tuple(context)).unhash(c)

    def valid_hash(self, context: List[N], c: str):
        return AlphabeticNameHash(tuple(context)).valid_hash(c)

    def parse(self, context: List[N], decoded_out: str,
              currently_selected: List[N]) -> List[str]:
        h = AlphabeticNameHash(tuple(context))
        logging.debug(f"{context=}")
        # do inference
        logging.debug(f"{decoded_out=}")
        selection, deselection = self.parse_raw(decoded_out)

        hashed_currently_selected = {h.hash(n) for n in currently_selected}
        desel_to_remove = deselection - hashed_currently_selected
        if len(desel_to_remove) > 0:
            logging.debug(f"warn! {desel_to_remove=}")
            deselection = deselection - desel_to_remove

        sel_to_remove = selection & hashed_currently_selected
        if len(sel_to_remove) > 0:
            logging.debug(f"warn! {sel_to_remove=}")
            selection = selection - sel_to_remove

        logging.debug("post strict cleaning")
        logging.debug(f"{selection=}")
        logging.debug(f"{deselection=}")

        model_clicks = selection | deselection
        logging.debug(f"{model_clicks=}")
        model_clicks_png = [h.unhash(n)
                            for n in model_clicks if h.valid_hash(n)]
        logging.debug(f"{model_clicks_png=}")
        return model_clicks_png

    @staticmethod
    def parse_raw(text: str) -> Tuple[Set[N], Set[N]]:
        last_answer = text.strip()
        if ":" in text:
            last_answer_pattern = r":.*$"
            xs = re.findall(last_answer_pattern, text)
            last_answer = xs[0].removeprefix(":").strip()
        xs = re.search(RE_PATTERN, last_answer)
        if xs is None:
            print(f"{last_answer=}")
            print("did not pass regex")
            return set(), set()

        select_pattern = r"(?<!de)select( [A-J])+$"
        xs = re.search(select_pattern, last_answer)
        if xs is not None:
            xs = xs.group()
        selections = set(xs.split(" ")[1:]) if xs else set()

        deselect_pattern = r"^deselect( [A-J])+"
        xs = re.search(deselect_pattern, last_answer)
        if xs is not None:
            xs = xs.group()
        deselections = set(xs.split(" ")[1:]) if xs else set()

        return selections, deselections

    def compose(self, context, chats, previous_selected, hash_images, padding):
        select_accum, deselect_accum, clickss = self.unfold_select_deselect(
            previous_selected)

        select_accum = select_accum + [[]]
        deselect_accum = deselect_accum + [[]]
        previous_selected = [[]] + previous_selected  # old states pre click
        assert len(chats) == len(select_accum) == len(
            deselect_accum) == len(previous_selected)

        messages, images = self.build_processor_input(
            context, chats, select_accum, deselect_accum, previous_selected, hash_images, omit_last_answer=True, sort_names=True, omit_context=False, chat_feedback=None)
        prompt = self.processor.apply_chat_template(
            messages, add_generation_prompt=True)
        prompt = prompt.strip()
        logging.debug(prompt)
        # Keep consistent with train_script
        inputs = self.processor(
            text=prompt, images=images,
            padding=padding, truncation=True, max_length=self.t_max_length,
            return_tensors="pt")
        return inputs

    def build_processor_input(self, image_pngs: List[N], chats: List[str],
                              select_accum: List[List[N]],
                              deselect_accum: List[List[N]],
                              pre_click_selected_accum: List[List[N]],
                              hash_image: bool, omit_last_answer: bool,
                              sort_names: bool, omit_context: bool,
                              chat_feedback: str, ):
        def _text_content(text): return {"type": "text", "text": text}

        def _image_content(): return {"type": "image"}

        def _user_prompt(content): return {"role": "user", "content": content}

        def _assistant_prompt(content): return {
            "role": "assistant", "content": content}

        def _system_prompt(content): return {
            "role": "system", "content": content}

        def _current_state(selected: List[N]):
            if len(selected) == 0:
                return 'none is selected'
            return f'{" ".join(selected)} currently selected'

        def _listener_action(select: List[N], deselect: List[N]):
            if len(select) == 0 and len(deselect) == 0:
                return 'nothing'
            if len(select) == 0:
                return f'deselect {" ".join(deselect)}'
            if len(deselect) == 0:
                return f'select {" ".join(select)}'
            return f'deselect {" ".join(deselect)} select {" ".join(select)}'

        func = AlphabeticNameHash(tuple(image_pngs)).hash if hash_image else id
        context, select_accum, deselect_accum,  pre_click_selected_accum = nested_apply(
            func, (image_pngs, select_accum, deselect_accum, pre_click_selected_accum))

        prompt = []
        images = []
        if not omit_context:
            images = [self.get_image(im) for im in image_pngs]
            images_and_names_content = []
            for im_name in context:
                images_and_names_content.append(_image_content())
                images_and_names_content.append(_text_content(im_name))
            prompt.append(_system_prompt(images_and_names_content))
        if not len(chats) == len(select_accum) == len(deselect_accum) == len(pre_click_selected_accum):
            logging.error(f"{chats=}")
            logging.error(f"{select_accum=}")
            logging.error(f"{deselect_accum=}")
            logging.error(f"{pre_click_selected_accum=}")
            assert False
        for i, (chat, select, deselect, pre_click_selected) in enumerate(
            zip(chats, select_accum, deselect_accum, pre_click_selected_accum)):
            if sort_names:
                select = sorted(select)
                deselect = sorted(deselect)
                pre_click_selected = sorted(pre_click_selected)

            prompt.append(_system_prompt(
                [_text_content(_current_state(pre_click_selected))]))
            prompt.append(_user_prompt([_text_content(chat)]))
            prompt.append(_assistant_prompt(
                [_text_content(_listener_action(select, deselect))]))
        if omit_last_answer:
            # idefics2 has processor.apply_chat_template(messages, add_generation_prompt=True) instead
            prompt.pop(-1)
        if chat_feedback is not None:
            prompt.append(_user_prompt([_text_content(chat_feedback)]))
        return prompt, images

    def unfold_select_deselect(self, previous_selected: List[List[N]]) -> Tuple[List[N], List[N], List[N]]:
        # currently selected AFTER i-th turn
        num_turns = len(previous_selected)
        selected: List[List[str]] = []  # turn-wise selection
        deselected: List[List[str]] = []  # turn-wise deselection
        clicks: List[List[str]] = []
        # combining turn-wise newly selected and newly deselected
        prev_selected = set()
        for turn in range(num_turns):
            curr_selected = set(previous_selected[turn])
            newly_selected = curr_selected - prev_selected
            newly_deselected = prev_selected - curr_selected
            selected.append(sorted_list(newly_selected))
            deselected.append(sorted_list(newly_deselected))
            clicks.append(sorted_list(newly_selected | newly_deselected))
            prev_selected = curr_selected.copy()
        return selected, deselected, clicks