import json
from abc import ABC
from typing import List
from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader


class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age


class Dialogue:
    """
    Build an abstract dialogue model using classes and methods to represent different dialogue elements.
    This class serves as a fundamental framework for constructing dialogue models.
    """

    def __init__(self, file_path: str):
        self.file_path = file_path
        self.turns = []

    def add_turn(self, turn):
        """
        Create an instance of a conversation participant
        :param turn:
        :return:
        """
        self.turns.append(turn)

    def parse_dialogue(self):
        """
        The parse_dialogue function reads the specified dialogue file and parses each dialogue turn line by line.
        For each turn, the function extracts the name of the speaker and the message content from the text,
        creating a Turn instance. If the speaker is not already present in the participants dictionary,
        a new Person instance is created. Finally, the parsed Turn instance is added to the Dialogue object.

        Please note that this sample code assumes that each line in the file follows a specific format:
        <speaker>:\r\n<message>\r\n\r\n. If your file has a different format or includes other metadata,
         you may need to adjust the parsing logic accordingly.
        """
        participants = {}
        speaker_name = None
        message = None

        with open(self.file_path, encoding='utf-8') as file:
            lines = file.readlines()
            for i, line in enumerate(lines):
                line = line.strip()
                if not line:
                    continue

                if speaker_name is None:
                    speaker_name, _ = line.split(':', 1)
                elif message is None:
                    message = line
                    if speaker_name not in participants:
                        participants[speaker_name] = Person(speaker_name, None)

                    speaker = participants[speaker_name]
                    turn = Turn(speaker, message)
                    self.add_turn(turn)

                    # Reset speaker_name and message for the next turn
                    speaker_name = None
                    message = None

    def display(self):
        for turn in self.turns:
            print(f"{turn.speaker.name}: {turn.message}")

    def export_to_file(self, file_path):
        with open(file_path, 'w', encoding='utf-8') as file:
            for turn in self.turns:
                file.write(f"{turn.speaker.name}: {turn.message}\n")

    def to_dict(self):
        dialogue_dict = {"turns": []}
        for turn in self.turns:
            turn_dict = {
                "speaker": turn.speaker.name,
                "message": turn.message
            }
            dialogue_dict["turns"].append(turn_dict)
        return dialogue_dict

    def to_json(self):
        dialogue_dict = self.to_dict()
        return json.dumps(dialogue_dict, ensure_ascii=False, indent=2)

    def participants_to_export(self):
        """
        participants_to_export
        :return:
        """
        participants = set()
        for turn in self.turns:
            participants.add(turn.speaker.name)
        return ', '.join(participants)


class Turn:
    def __init__(self, speaker, message):
        self.speaker = speaker
        self.message = message


class DialogueLoader(BaseLoader, ABC):
    """Load dialogue."""

    def __init__(self, file_path: str):
        """Initialize with dialogue."""
        self.file_path = file_path
        dialogue = Dialogue(file_path=file_path)
        dialogue.parse_dialogue()
        self.dialogue = dialogue

    def load(self) -> List[Document]:
        """Load from dialogue."""
        documents = []
        participants = self.dialogue.participants_to_export()

        for turn in self.dialogue.turns:
            metadata = {"source": f"Dialogue File：{self.dialogue.file_path},"
                                  f"speaker：{turn.speaker.name}，"
                                  f"participant：{participants}"}
            turn_document = Document(page_content=turn.message, metadata=metadata.copy())
            documents.append(turn_document)

        return documents