from dataclasses import dataclass
import logging
import os
from abc import ABC
from typing import Optional

import torch
import json

from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
)
from ts.torch_handler.base_handler import BaseHandler

logger = logging.getLogger(__name__)

MAX_TOKEN_LENGTH_ERR = {
    "code": 422,
    "type" : "MaxTokenLengthError",
    "message": "Max token length exceeded",
}

class CopEngHandler(BaseHandler, ABC):

    @dataclass
    class GenerationConfig:
        max_length: int = 20
        max_new_tokens: Optional[int] = None
        min_length: int = 0
        min_new_tokens: Optional[int] = None
        early_stopping: bool = True
        do_sample: bool = False
        num_beams: int = 1
        num_beam_groups: int = 1
        top_k: int = 50
        top_p: float = 0.95
        temperature: float = 1.0
        diversity_penalty: float = 0.0

    def __init__(self):
        super(CopEngHandler, self).__init__()
        self.initialized = False

    def initialize(self, ctx):
        """In this initialize function, the HF large model is loaded and
        partitioned using DeepSpeed.
        Args:
            ctx (context): It is a JSON Object containing information
            pertaining to the model artifacts parameters.
        """
        logger.info("Start initialize")
        self.manifest = ctx.manifest
        properties = ctx.system_properties
        model_dir = properties.get("model_dir")
        serialized_file = self.manifest["model"]["serializedFile"]
        model_pt_path = os.path.join(model_dir, serialized_file)

        setup_config_path = os.path.join(model_dir, "setup_self.config.json")
        if os.path.isfile(setup_config_path):
            with open(setup_config_path) as setup_config_path:
                self.setup_config = json.load(setup_config_path)

        seed = int(42)
        torch.manual_seed(seed)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info("Device: %s", self.device)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
        self.model.to(self.device)
        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.config = CopEngHandler.GenerationConfig(
            max_new_tokens=128,
            min_new_tokens=1,
            num_beams=5,
        )
        self.initialized = True
        logger.info("Init done")
                


    def preprocess(self, requests):
        preprocessed_data = []
        for data in requests:
            data_item = data.get("data")
            if data_item is None:
                data_item = data.get("body")
            if isinstance(data_item, (bytes, bytearray)):
                data_item = data_item.decode("utf-8")
            preprocessed_data.append(greekify(data_item))
        logger.info("preprocessed_data %s: ", preprocessed_data)
        return preprocessed_data

    def inference(self, data):
        indices = {}
        batch = []
        for i, item in enumerate(data):
            tokens = self.tokenizer(item, return_tensors="pt", padding=True)
            if len(tokens.input_ids.squeeze()) > self.tokenizer.model_max_length:
                logger.info("Skipping token %s for index %s", tokens, i)
                continue
            indices[i] = len(batch)
            batch.append(data[i])
        logger.info("inference batch: %s", batch)
        result = self.batch_translate(batch)
        return [result[indices[i]] if i in indices else None for i in range(len(data))]

    def postprocess(self, output):
        return output

    def handle(self, requests, context):
        preprocessed = self.preprocess(requests)
        inference_data = self.inference(preprocessed)
        postprocessed = self.postprocess(inference_data)
        logger.info("inference result: %s", postprocessed)

        responses = [
            {"code": 200, "translation": translation}
            if translation
            else MAX_TOKEN_LENGTH_ERR
            for translation in postprocessed
        ]

        return responses

    def batch_translate(self, input_sentences, output_confidence=False):
        if len(input_sentences) == 0:
            return []
        inputs = self.tokenizer(input_sentences, return_tensors="pt", padding=True).to(
            self.device
        )
        output_scores, return_dict_in_generate = output_confidence, output_confidence
        outputs = self.model.generate(
            **inputs,
            max_length=self.config.max_length,
            max_new_tokens=self.config.max_new_tokens,
            min_length=self.config.min_length,
            min_new_tokens=self.config.min_new_tokens,
            early_stopping=self.config.early_stopping,
            do_sample=self.config.do_sample,
            num_beams=self.config.num_beams,
            num_beam_groups=self.config.num_beam_groups,
            top_k=self.config.top_k,
            top_p=self.config.top_p,
            temperature=self.config.temperature,
            diversity_penalty=self.config.diversity_penalty,
            output_scores=output_scores,
            return_dict_in_generate=True,
        )
        translated_text = self.tokenizer.batch_decode(
            outputs.sequences, skip_special_tokens=True
        )
        return translated_text


COPTIC_TO_GREEK = {
    "ⲁ": "α",
    "ⲃ": "β",
    "ⲅ": "γ",
    "ⲇ": "δ",
    "ⲉ": "ε",
    "ⲋ": "ϛ",
    "ⲍ": "ζ",
    "ⲏ": "η",
    "ⲑ": "θ",
    "ⲓ": "ι",
    "ⲕ": "κ",
    "ⲗ": "λ",
    "ⲙ": "μ",
    "ⲛ": "ν",
    "ⲝ": "ξ",
    "ⲟ": "ο",
    "ⲡ": "π",
    "ⲣ": "ρ",
    "ⲥ": "σ",
    "ⲧ": "τ",
    "ⲩ": "υ",
    "ⲫ": "φ",
    "ⲭ": "χ",
    "ⲯ": "ψ",
    "ⲱ": "ω",
    "ϣ": "s",
    "ϥ": "f",
    "ϧ": "k",
    "ϩ": "h",
    "ϫ": "j",
    "ϭ": "c",
    "ϯ": "t",
}


def greekify(coptic_text):
    chars = []
    for c in coptic_text:
        l_c = c.lower()
        chars.append(COPTIC_TO_GREEK.get(l_c, l_c))
    return "".join(chars)