import os

import agentspeak
from agentspeak import AslError
import litellm

from a2a_agentspeak import actions
from a2a_acl.utils.strings import neutralize_str, clear

llm_model = "mistral/mistral-small-latest"

os.environ["LANGFUSE_OTEL_HOST"] = "https://cloud.langfuse.com"
litellm.callbacks = ["langfuse_otel"]


failure = agentspeak.Literal("llm_failure")


def log(m):
    print("[LOG] " + m)



def ask_llm_for_completion(spec: str, req_list: str, rejected_list = None):
    try:
        content = "Specification: " + spec + "(end of the specification)" + " List of requirements: " + req_list + "(end of list of requirements)." + ("" if rejected_list is None else (
                    "(Note that the following requirements do not apply: " + str(rejected_list) + ")"))
        chat_response = litellm.completion(
            model=llm_model,
            messages=[
                {
                    "role": "system",
                    "content": "Given a specification of a system, and a list of atomic requirements, give an atomic requirements that covers the specification and which is not included in the given list of requirements."
                    + " Answer with the new requirement, don't explain.",
                },
                {
                    "role": "user",
                    "content": content,
                },
            ],
            max_tokens=50,
            temperature=1.2,
        )
        res = chat_response.choices[0].message.content
        log(
            "I had an interaction with " + llm_model + " to generate a requirement.\n "
            + "I gave the following info: "
            + content
            +"\n"
            + "The LLM " + llm_model + " gave me the following answer: "
            + res
        )
        return res
    except Exception as e:
        print(str(e))
        raise AslError("LLM failure: " + str(type(e))) from e





def prompt_generation(s, req, rej) -> agentspeak.Literal:
    assert s.functor == "spec"
    assert req.functor == "req"
    assert rej.functor == "rejected_list"
    try:
        res = ask_llm_for_completion(str(s.args[0]), str(req.args[0]), str(rej.args[0]))
        res = clear(res)
        return agentspeak.Literal("generated_req", (res,))
    except Exception as e:
        print("Failure while talking with the llm")
        print(str(e))
        return failure


action_prompt_generate = actions.Action(
    "function",
    ".prompt_generate_alt",
    (agentspeak.Literal, agentspeak.Literal, agentspeak.Literal),
    prompt_generation,
)
