import subprocess

class DLVHandler:
    def __init__(self, filename="program.dlv", useDlv2=False):
        self.filename = filename
        self.program_content = ""
        self.useDlv2 = useDlv2

    def add_fact(self, fact):
        """ 添加一个事实到程序中 """
        self.program_content += fact + ".\n"

    def add_facts(self, facts):
        """ 添加多个事实到程序中 """
        for fact in facts:
            self.add_fact(fact)

    def add_rule(self, rule):
        """ 添加一个规则到程序中 """
        self.program_content += rule + ".\n"

    def add_rules(self, rules):
        """ 添加多个规则到程序中 """
        for rule in rules:
            self.add_rule(rule)

    def add_query(self, query):
        """ 添加一个查询到程序中 """
        self.program_content += "query(" + query + ").\n"

    def clear_program(self):
        """ 清空当前程序内容 """
        self.program_content = ""

    def save_program(self):
        """ 将当前程序保存到文件中 """
        with open(self.filename, "w") as file:
            file.write(self.program_content)

    def run(self):
        """ 运行当前的 DLV 程序文件，并返回输出 """
        try:
            if self.useDlv2:
                # comment: 记得删掉
                result = subprocess.run(["dlv2", self.filename, "--no-facts"], text=True, capture_output=True, check=True)
            else:
                result = subprocess.run(["dlv", self.filename], text=True, capture_output=True, check=True)
            return result.stdout
        except subprocess.CalledProcessError as e:
            return f"Error during DLV execution: {e}"

    def run_and_get_results(self):
        """ 运行当前的 DLV 程序文件，并返回查询结果 """
        self.save_program()
        output = self.run()
        self.clear_program()
        return [line.strip() for line in output.split('\n') if line.strip()][-1] # 返回查询结果

    def run_given_program(self, program_content):
        """ 运行给定的 DLV 程序内容，并返回输出 """
        self.program_content = program_content
        self.save_program()
        result = self.run()
        self.clear_program()
        return result

    def quick_run(self, facts, rules):
        """ 快速运行给定的 DLV 程序内容，并返回输出 """
        self.add_facts(facts)
        self.add_rules(rules)
        result = self.run_and_get_results()
        return result

# 使用示例
if __name__ == "__main__":
    handler = DLVHandler()
    handler.add_fact('bird("tweety")')
    handler.add_fact('bird("pingu")')
    handler.add_fact('penguin("pingu")')
    handler.add_rule('flies(X) :- bird(X), not penguin(X)')
    handler.add_query('flies("pingu")')
    handler.save_program()
    output = handler.run()
    print("DLV Output:")
    print(output)
