import uuid
from collections.abc import Iterable

import zmq
from pydantic import ValidationError
from slugify import slugify
from zmq.asyncio import Context

from control_backend.agents import BaseAgent
from control_backend.core.config import settings
from control_backend.schemas.program import (
    Action,
    BasicBelief,
    BasicNorm,
    Belief,
    ConditionalNorm,
    GestureAction,
    Goal,
    InferredBelief,
    KeywordBelief,
    LLMAction,
    LogicalOperator,
    Phase,
    Plan,
    Program,
    ProgramElement,
    SemanticBelief,
    SpeechAction,
    Trigger,
)

test_program = Program(
    phases=[
        Phase(
            norms=[
                BasicNorm(norm="Talk like a pirate", id=uuid.uuid4()),
                ConditionalNorm(
                    condition=InferredBelief(
                        left=KeywordBelief(keyword="Arr", id=uuid.uuid4()),
                        right=SemanticBelief(
                            description="testing", name="semantic belief", id=uuid.uuid4()
                        ),
                        operator=LogicalOperator.OR,
                        name="Talking to a pirate",
                        id=uuid.uuid4(),
                    ),
                    norm="Use nautical terms",
                    id=uuid.uuid4(),
                ),
                ConditionalNorm(
                    condition=SemanticBelief(
                        description="We are talking to a child",
                        name="talking to child",
                        id=uuid.uuid4(),
                    ),
                    norm="Do not use cuss words",
                    id=uuid.uuid4(),
                ),
            ],
            triggers=[
                Trigger(
                    condition=InferredBelief(
                        left=KeywordBelief(keyword="key", id=uuid.uuid4()),
                        right=InferredBelief(
                            left=KeywordBelief(keyword="key2", id=uuid.uuid4()),
                            right=SemanticBelief(
                                description="Decode this", name="semantic belief 2", id=uuid.uuid4()
                            ),
                            operator=LogicalOperator.OR,
                            name="test trigger inferred inner",
                            id=uuid.uuid4(),
                        ),
                        operator=LogicalOperator.OR,
                        name="test trigger inferred outer",
                        id=uuid.uuid4(),
                    ),
                    plan=Plan(
                        steps=[
                            SpeechAction(text="Testing trigger", id=uuid.uuid4()),
                            Goal(
                                name="Testing trigger",
                                plan=Plan(
                                    steps=[LLMAction(goal="Do something", id=uuid.uuid4())],
                                    id=uuid.uuid4(),
                                ),
                                id=uuid.uuid4(),
                            ),
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                )
            ],
            goals=[
                Goal(
                    name="Determine user age",
                    plan=Plan(
                        steps=[LLMAction(goal="Determine the age of the user.", id=uuid.uuid4())],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
                Goal(
                    name="Find the user's name",
                    plan=Plan(
                        steps=[
                            Goal(
                                name="Greet the user",
                                plan=Plan(
                                    steps=[LLMAction(goal="Greet the user.", id=uuid.uuid4())],
                                    id=uuid.uuid4(),
                                ),
                                can_fail=False,
                                id=uuid.uuid4(),
                            ),
                            Goal(
                                name="Ask for name",
                                plan=Plan(
                                    steps=[
                                        LLMAction(goal="Obtain the user's name.", id=uuid.uuid4())
                                    ],
                                    id=uuid.uuid4(),
                                ),
                                id=uuid.uuid4(),
                            ),
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
                Goal(
                    name="Tell a joke",
                    plan=Plan(
                        steps=[LLMAction(goal="Tell a joke.", id=uuid.uuid4())], id=uuid.uuid4()
                    ),
                    id=uuid.uuid4(),
                ),
            ],
            id=uuid.uuid4(),
        ),
        Phase(
            id=uuid.uuid4(),
            norms=[
                BasicNorm(norm="Use very gentle speech.", id=uuid.uuid4()),
                ConditionalNorm(
                    condition=SemanticBelief(
                        description="We are talking to a child",
                        name="talking to child",
                        id=uuid.uuid4(),
                    ),
                    norm="Do not use cuss words",
                    id=uuid.uuid4(),
                ),
            ],
            triggers=[
                Trigger(
                    condition=InferredBelief(
                        left=KeywordBelief(keyword="help", id=uuid.uuid4()),
                        right=SemanticBelief(
                            description="User is stuck", name="stuck", id=uuid.uuid4()
                        ),
                        operator=LogicalOperator.OR,
                        name="help_or_stuck",
                        id=uuid.uuid4(),
                    ),
                    plan=Plan(
                        steps=[
                            Goal(
                                name="Unblock user",
                                plan=Plan(
                                    steps=[
                                        LLMAction(
                                            goal="Provide a step-by-step path to "
                                            "resolve the user's issue.",
                                            id=uuid.uuid4(),
                                        )
                                    ],
                                    id=uuid.uuid4(),
                                ),
                                id=uuid.uuid4(),
                            ),
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
            ],
            goals=[
                Goal(
                    name="Clarify intent",
                    plan=Plan(
                        steps=[
                            LLMAction(
                                goal="Ask 1-2 targeted questions to clarify the "
                                "user's intent, then proceed.",
                                id=uuid.uuid4(),
                            )
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
                Goal(
                    name="Provide solution",
                    plan=Plan(
                        steps=[
                            LLMAction(
                                goal="Deliver a solution to complete the user's goal.",
                                id=uuid.uuid4(),
                            )
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
                Goal(
                    name="Summarize next steps",
                    plan=Plan(
                        steps=[
                            LLMAction(
                                goal="Summarize what the user should do next.", id=uuid.uuid4()
                            )
                        ],
                        id=uuid.uuid4(),
                    ),
                    id=uuid.uuid4(),
                ),
            ],
        ),
    ]
)


def do_things():
    print(AgentSpeakGenerator().generate(test_program))


class AgentSpeakGenerator:
    """
    Converts Pydantic representation of behavior programs into AgentSpeak(L) code string.
    """

    arrow_prefix = f"{' ' * 2}<-{' ' * 2}"
    colon_prefix = f"{' ' * 2}:{' ' * 3}"
    indent_prefix = " " * 6

    def generate(self, program: Program) -> str:
        lines = []
        lines.append("")

        lines += self._generate_initial_beliefs(program)

        lines += self._generate_basic_flow(program)

        lines += self._generate_phase_transitions(program)

        lines += self._generate_norms(program)

        lines += self._generate_belief_inference(program)

        lines += self._generate_goals(program)

        lines += self._generate_triggers(program)

        return "\n".join(lines)

    def _generate_initial_beliefs(self, program: Program) -> Iterable[str]:
        yield "// --- Initial beliefs and agent startup ---"

        yield "phase(start)."

        yield ""

        yield "+started"
        yield f"{self.colon_prefix}phase(start)"
        yield f"{self.arrow_prefix}phase({program.phases[0].id if program.phases else 'end'})."

        yield from ["", ""]

    def _generate_basic_flow(self, program: Program) -> Iterable[str]:
        yield "// --- Basic flow ---"

        for phase in program.phases:
            yield from self._generate_basic_flow_per_phase(phase)

        yield from ["", ""]

    def _generate_basic_flow_per_phase(self, phase: Phase) -> Iterable[str]:
        yield "+user_said(Message)"
        yield f"{self.colon_prefix}phase({phase.id})"

        goals = phase.goals
        if goals:
            yield f"{self.arrow_prefix}{self._slugify(goals[0], include_prefix=True)}"
            for goal in goals[1:]:
                yield f"{self.indent_prefix}{self._slugify(goal, include_prefix=True)}"

        yield f"{self.indent_prefix if goals else self.arrow_prefix}!transition_phase."

    def _generate_phase_transitions(self, program: Program) -> Iterable[str]:
        yield "// --- Phase transitions ---"

        if len(program.phases) == 0:
            yield from ["", ""]
            return

        # TODO: remove outdated things

        for i in range(-1, len(program.phases)):
            predecessor = program.phases[i] if i >= 0 else None
            successor = program.phases[i + 1] if i < len(program.phases) - 1 else None
            yield from self._generate_phase_transition(predecessor, successor)

        yield from self._generate_phase_transition(None, None)  # to avoid failing plan

        yield from ["", ""]

    def _generate_phase_transition(
        self, phase: Phase | None = None, next_phase: Phase | None = None
    ) -> Iterable[str]:
        yield "+!transition_phase"

        if phase is None and next_phase is None:  # base case true to avoid failing plan
            yield f"{self.arrow_prefix}true."
            return

        yield f"{self.colon_prefix}phase({phase.id if phase else 'start'})"
        yield f"{self.arrow_prefix}-+phase({next_phase.id if next_phase else 'end'})."

    def _generate_norms(self, program: Program) -> Iterable[str]:
        yield "// --- Norms ---"

        for phase in program.phases:
            for norm in phase.norms:
                if type(norm) is BasicNorm:
                    yield f"{self._slugify(norm)} :- phase({phase.id})."
                if type(norm) is ConditionalNorm:
                    yield (
                        f"{self._slugify(norm)} :- phase({phase.id}) & "
                        f"{self._slugify(norm.condition)}."
                    )

        yield from ["", ""]

    def _generate_belief_inference(self, program: Program) -> Iterable[str]:
        yield "// --- Belief inference rules ---"

        for phase in program.phases:
            for norm in phase.norms:
                if not isinstance(norm, ConditionalNorm):
                    continue

                yield from self._belief_inference_recursive(norm.condition)

            for trigger in phase.triggers:
                yield from self._belief_inference_recursive(trigger.condition)

        yield from ["", ""]

    def _belief_inference_recursive(self, belief: Belief) -> Iterable[str]:
        if type(belief) is KeywordBelief:
            yield (
                f"{self._slugify(belief)} :- user_said(Message) & "
                f'.substring(Message, "{belief.keyword}", Pos) & Pos >= 0.'
            )
        if type(belief) is InferredBelief:
            yield (
                f"{self._slugify(belief)} :- {self._slugify(belief.left)} "
                f"{'&' if belief.operator == LogicalOperator.AND else '|'} "
                f"{self._slugify(belief.right)}."
            )

            yield from self._belief_inference_recursive(belief.left)
            yield from self._belief_inference_recursive(belief.right)

    def _generate_goals(self, program: Program) -> Iterable[str]:
        yield "// --- Goals ---"

        for phase in program.phases:
            previous_goal: Goal | None = None
            for goal in phase.goals:
                yield from self._generate_goal_plan_recursive(goal, phase, previous_goal)
                previous_goal = goal

        yield from ["", ""]

    def _generate_goal_plan_recursive(
        self, goal: Goal, phase: Phase, previous_goal: Goal | None = None
    ) -> Iterable[str]:
        yield f"+{self._slugify(goal, include_prefix=True)}"

        # Context
        yield f"{self.colon_prefix}phase({phase.id}) &"
        yield f"{self.indent_prefix}not responded_this_turn &"
        yield f"{self.indent_prefix}not achieved_{self._slugify(goal)} &"
        if previous_goal:
            yield f"{self.indent_prefix}achieved_{self._slugify(previous_goal)}"
        else:
            yield f"{self.indent_prefix}true"

        extra_goals_to_generate = []

        steps = goal.plan.steps

        if len(steps) == 0:
            yield f"{self.arrow_prefix}true."
            return

        first_step = steps[0]
        yield (
            f"{self.arrow_prefix}{self._slugify(first_step, include_prefix=True)}"
            f"{'.' if len(steps) == 1 and goal.can_fail else ';'}"
        )
        if isinstance(first_step, Goal):
            extra_goals_to_generate.append(first_step)

        for step in steps[1:-1]:
            yield f"{self.indent_prefix}{self._slugify(step, include_prefix=True)};"
            if isinstance(step, Goal):
                extra_goals_to_generate.append(step)

        if len(steps) > 1:
            last_step = steps[-1]
            yield (
                f"{self.indent_prefix}{self._slugify(last_step, include_prefix=True)}"
                f"{'.' if goal.can_fail else ';'}"
            )
            if isinstance(last_step, Goal):
                extra_goals_to_generate.append(last_step)

        if not goal.can_fail:
            yield f"{self.indent_prefix}+achieved_{self._slugify(goal)}."

        yield f"+{self._slugify(goal, include_prefix=True)}"
        yield f"{self.arrow_prefix}true."

        yield ""

        extra_previous_goal: Goal | None = None
        for extra_goal in extra_goals_to_generate:
            yield from self._generate_goal_plan_recursive(extra_goal, phase, extra_previous_goal)
            extra_previous_goal = extra_goal

    def _generate_triggers(self, program: Program) -> Iterable[str]:
        yield "// --- Triggers ---"

        for phase in program.phases:
            for trigger in phase.triggers:
                yield from self._generate_trigger_plan(trigger, phase)

        yield from ["", ""]

    def _generate_trigger_plan(self, trigger: Trigger, phase: Phase) -> Iterable[str]:
        belief_name = self._slugify(trigger.condition)

        yield f"+{belief_name}"
        yield f"{self.colon_prefix}phase({phase.id})"

        extra_goals_to_generate = []

        steps = trigger.plan.steps

        if len(steps) == 0:
            yield f"{self.arrow_prefix}true."
            return

        first_step = steps[0]
        yield (
            f"{self.arrow_prefix}{self._slugify(first_step, include_prefix=True)}"
            f"{'.' if len(steps) == 1 else ';'}"
        )
        if isinstance(first_step, Goal):
            extra_goals_to_generate.append(first_step)

        for step in steps[1:-1]:
            yield f"{self.indent_prefix}{self._slugify(step, include_prefix=True)};"
            if isinstance(step, Goal):
                extra_goals_to_generate.append(step)

        if len(steps) > 1:
            last_step = steps[-1]
            yield f"{self.indent_prefix}{self._slugify(last_step, include_prefix=True)}."
            if isinstance(last_step, Goal):
                extra_goals_to_generate.append(last_step)

        yield ""

        extra_previous_goal: Goal | None = None
        for extra_goal in extra_goals_to_generate:
            yield from self._generate_trigger_plan_recursive(extra_goal, phase, extra_previous_goal)
            extra_previous_goal = extra_goal

    def _generate_trigger_plan_recursive(
        self, goal: Goal, phase: Phase, previous_goal: Goal | None = None
    ) -> Iterable[str]:
        yield f"+{self._slugify(goal, include_prefix=True)}"

        extra_goals_to_generate = []

        steps = goal.plan.steps

        if len(steps) == 0:
            yield f"{self.arrow_prefix}true."
            return

        first_step = steps[0]
        yield (
            f"{self.arrow_prefix}{self._slugify(first_step, include_prefix=True)}"
            f"{'.' if len(steps) == 1 and goal.can_fail else ';'}"
        )
        if isinstance(first_step, Goal):
            extra_goals_to_generate.append(first_step)

        for step in steps[1:-1]:
            yield f"{self.indent_prefix}{self._slugify(step, include_prefix=True)};"
            if isinstance(step, Goal):
                extra_goals_to_generate.append(step)

        if len(steps) > 1:
            last_step = steps[-1]
            yield (
                f"{self.indent_prefix}{self._slugify(last_step, include_prefix=True)}"
                f"{'.' if goal.can_fail else ';'}"
            )
            if isinstance(last_step, Goal):
                extra_goals_to_generate.append(last_step)

        if not goal.can_fail:
            yield f"{self.indent_prefix}+achieved_{self._slugify(goal)}."

        yield f"+{self._slugify(goal, include_prefix=True)}"
        yield f"{self.arrow_prefix}true."

        yield ""

        extra_previous_goal: Goal | None = None
        for extra_goal in extra_goals_to_generate:
            yield from self._generate_goal_plan_recursive(extra_goal, phase, extra_previous_goal)
            extra_previous_goal = extra_goal

    def _slugify(self, element: ProgramElement, include_prefix: bool = False) -> str:
        def base_slugify_call(text: str):
            return slugify(text, separator="_", stopwords=["a", "the"])

        if type(element) is KeywordBelief:
            return f'keyword_said("{element.keyword}")'

        if type(element) is SemanticBelief:
            name = element.name
            return f"semantic_{base_slugify_call(name if name else element.description)}"

        if isinstance(element, BasicNorm):
            return f'norm("{element.norm}")'

        if isinstance(element, Goal):
            return f"{'!' if include_prefix else ''}{base_slugify_call(element.name)}"

        if isinstance(element, SpeechAction):
            return f'.say("{element.text}")'

        if isinstance(element, GestureAction):
            return f'.gesture("{element.gesture}")'

        if isinstance(element, LLMAction):
            return f'!generate_response_with_goal("{element.goal}")'

        if isinstance(element, Action.__value__):
            raise NotImplementedError(
                "Have not implemented an ASL string representation for this action."
            )

        if element.name == "":
            raise ValueError("Name must be initialized for this type of ProgramElement.")

        return base_slugify_call(element.name)

    def _extract_basic_beliefs_from_program(self, program: Program) -> list[BasicBelief]:
        beliefs = []

        for phase in program.phases:
            for norm in phase.norms:
                if isinstance(norm, ConditionalNorm):
                    beliefs += self._extract_basic_beliefs_from_belief(norm.condition)

            for trigger in phase.triggers:
                beliefs += self._extract_basic_beliefs_from_belief(trigger.condition)

        return beliefs

    def _extract_basic_beliefs_from_belief(self, belief: Belief) -> list[BasicBelief]:
        if isinstance(belief, InferredBelief):
            return self._extract_basic_beliefs_from_belief(
                belief.left
            ) + self._extract_basic_beliefs_from_belief(belief.right)
        return [belief]


class BDIProgramManager(BaseAgent):
    """
    BDI Program Manager Agent.

    This agent is responsible for receiving high-level programs (sequences of instructions/goals)
    from the external HTTP API (via ZMQ) and translating them into core beliefs (norms and goals)
    for the BDI Core Agent. In the future, it will be responsible for determining when goals are
    met, and passing on new norms and goals accordingly.

    :ivar sub_socket: The ZMQ SUB socket used to receive program updates.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.sub_socket = None

    # async def _send_to_bdi(self, program: Program):
    #     """
    #     Convert a received program into BDI beliefs and send them to the BDI Core Agent.
    #
    #     Currently, it takes the **first phase** of the program and extracts:
    #     -   **Norms**: Constraints or rules the agent must follow.
    #     -   **Goals**: Objectives the agent must achieve.
    #
    #     These are sent as a ``BeliefMessage`` with ``replace=True``, meaning they will
    #     overwrite any existing norms/goals of the same name in the BDI agent.
    #
    #     :param program: The program object received from the API.
    #     """
    #     first_phase = program.phases[0]
    #     norms_belief = Belief(
    #         name="norms",
    #         arguments=[norm.norm for norm in first_phase.norms],
    #         replace=True,
    #     )
    #     goals_belief = Belief(
    #         name="goals",
    #         arguments=[goal.description for goal in first_phase.goals],
    #         replace=True,
    #     )
    #     program_beliefs = BeliefMessage(beliefs=[norms_belief, goals_belief])
    #
    #     message = InternalMessage(
    #         to=settings.agent_settings.bdi_core_name,
    #         sender=self.name,
    #         body=program_beliefs.model_dump_json(),
    #         thread="beliefs",
    #     )
    #     await self.send(message)
    #     self.logger.debug("Sent new norms and goals to the BDI agent.")

    async def _receive_programs(self):
        """
        Continuous loop that receives program updates from the HTTP endpoint.

        It listens to the ``program`` topic on the internal ZMQ SUB socket.
        When a program is received, it is validated and forwarded to BDI via :meth:`_send_to_bdi`.
        """
        while True:
            topic, body = await self.sub_socket.recv_multipart()

            try:
                program = Program.model_validate_json(body)
            except ValidationError:
                self.logger.exception("Received an invalid program.")
                continue

            await self._send_to_bdi(program)

    async def setup(self):
        """
        Initialize the agent.

        Connects the internal ZMQ SUB socket and subscribes to the 'program' topic.
        Starts the background behavior to receive programs.
        """
        context = Context.instance()

        self.sub_socket = context.socket(zmq.SUB)
        self.sub_socket.connect(settings.zmq_settings.internal_sub_address)
        self.sub_socket.subscribe("program")

        self.add_behavior(self._receive_programs())


if __name__ == "__main__":
    do_things()