import uuid
import json


class Turn:
    def __init__(self, query, response, turn_id=None, source_model="") -> None:
        if turn_id:
            self.turn_id = turn_id
        else:
            self.turn_id = uuid.uuid4().hex

        self.query = query
        self.response = response

        # generated by what model
        self.source_model = source_model

    def dumps(self):
        return {
            "turn_id": self.turn_id,
            "query": self.query,
            "response": self.response,
            "source_model": self.source_model,
        }

    def loads(self, json_data):
        if "turn_id" in json_data:
            self.turn_id = json_data["turn_id"]
        self.query = json_data["query"]
        self.response = json_data["response"]
        self.source_model = json_data["source_model"]


class PairwiseTurn(Turn):
    def __init__(self, query, pairwise_response, turn_id=None, source_model="") -> None:
        super().__init__(query, None, turn_id, source_model)
        self.pairwise_response = pairwise_response

    def dumps(self):
        dump_json = super().dumps()
        dump_json["pairwise_response"] = self.pairwise_response
        return dump_json

    def loads(self, json_data):
        super().loads(json_data)
        self.pairwise_response = json_data["pairwise_response"]


class ReferenceTurn(Turn):
    def __init__(
        self, query, response, reference, turn_id=None, source_model=""
    ) -> None:
        super().__init__(query, response, turn_id, source_model)
        self.reference = reference

    def dumps(self):
        dump_json = super().dumps()
        dump_json["reference"] = self.reference
        return dump_json

    def loads(self, json_data):
        super().loads(json_data)
        self.reference = json_data["reference"]


class Dialogue:
    def __init__(self, session_id=None) -> None:
        if session_id:
            self.session_id = session_id
        else:
            self.session_id = uuid.uuid4().hex

        self.meta = {}
        self.turns = []
        self.meta_type = "Dialogue"

    def size(self):
        return len(self.turns)

    def add_turn_item(self, query, response, turn_id=None, source_model=""):
        turn = Turn(query, response, turn_id, source_model)
        self.turns.append(turn)

    def get_query_by_idx(self, turn_idx):
        return self.turns[turn_idx].query

    def set_query_by_idx(self, turn_idx, new_query):
        self.turns[turn_idx].query = new_query

    def get_response_by_idx(self, turn_idx):
        return self.turns[turn_idx].response

    def set_response_by_idx(self, turn_idx, new_response):
        self.turns[turn_idx].response = new_response

    def get_session_id(self):
        return self.session_id

    def set_meta_info(self, meta_info):
        self.meta = meta_info

    def get_meta_info(self):
        return self.meta

    def iter(self):
        for turn in self.turns:
            yield turn

    def dumps(self):
        data = {
            "session_id": self.session_id,
            "meta": self.meta,
            "meta_type": self.meta_type,
            "turns": [t.dumps() for t in self.turns],
        }
        return json.dumps(data, ensure_ascii=False)

    def loads(self, json_data):
        self.session_id = json_data["session_id"]
        self.meta = json_data["meta"]
        self.meta_type = json_data["meta_type"]
        for json_turn in json_data["turns"]:
            turn = Turn("", "")
            turn.loads(json_turn)
            self.turns.append(turn)


class ReferenceDialogue(Dialogue):
    def __init__(self, session_id=None) -> None:
        super().__init__(session_id)
        self.meta_type = "ReferenceDialogue"

    def add_turn_item(self, query, response, reference, turn_id=None, source_model=""):
        turn = ReferenceTurn(query, response, reference, turn_id, source_model)
        self.turns.append(turn)

    def dumps(self):
        data = {
            "session_id": self.session_id,
            "meta": self.meta,
            "turns": [t.dumps() for t in self.turns],
            "meta_type": self.meta_type,
        }
        return json.dumps(data, ensure_ascii=False)

    def loads(self, json_data):
        self.session_id = json_data["session_id"]
        self.meta = json_data["meta"]
        self.meta_type = json_data["meta_type"]
        for json_turn in json_data["turns"]:
            turn = ReferenceTurn("", "", "")
            turn.loads(json_turn)
            self.turns.append(turn)


class PairwiseDialogue(Dialogue):
    def __init__(self, session_id=None) -> None:
        super().__init__(session_id)
        self.meta_type = "PairwiseDialogue"

    def add_turn_item(self, query, pairwise_response, turn_id=None, source_model=""):
        turn = PairwiseTurn(query, pairwise_response, turn_id, source_model)
        self.turns.append(turn)

    def get_pairwise_key(self):
        return self.turns[0].pairwise_response.keys()

    def get_query_by_idx(self, turn_idx):
        return self.turns[turn_idx].query

    def get_pairwise_response_by_idx(self, turn_idx, key):
        return self.turns[turn_idx].pairwise_response[key]

    def dumps(self):
        data = {
            "session_id": self.session_id,
            "meta": self.meta,
            "turns": [t.dumps() for t in self.turns],
            "meta_type": self.meta_type,
        }
        return json.dumps(data, ensure_ascii=False)

    def loads(self, json_data):
        self.session_id = json_data["session_id"]
        self.meta = json_data["meta"]
        self.meta_type = json_data["meta_type"]
        for json_turn in json_data["turns"]:
            turn = PairwiseTurn("", "")
            turn.loads(json_turn)
            self.turns.append(turn)
