"""
Node definitions for the graph.
Includes base Node, ReflectionNode, RewritingNode, InputNode, and OutputNode.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict

from .llm_integration import get_llm


class BaseNode(ABC):
    """
    Abstract base class for all nodes in the graph.
    Each node must implement the `process` method.
    """

    def __init__(self, node_id: str):
        self.node_id = node_id

    @abstractmethod
    def process(self, input_data: Any) -> Any:
        """
        Process the input data and return the output.
        """
        pass


class InputNode(BaseNode):
    """
    Node that simply passes through the input data.
    """

    def process(self, input_data: Any) -> Any:
        return input_data


class OutputNode(BaseNode):
    """
    Node that collects the final output.
    """

    def process(self, input_data: Any) -> Any:
        return input_data


class ReflectionNode(BaseNode):
    """
    Node that generates reflective insights from the input text using an LLM.
    """

    def __init__(self, node_id: str, prompt_template: str = None):
        super().__init__(node_id)
        self.prompt_template = (
            prompt_template
            or "Please reflect on the following text:\n\n{input_text}\n\nReflection:"
        )
        self.llm = get_llm()

    def process(self, input_data: str) -> Dict[str, str]:
        prompt = self.prompt_template.format(input_text=input_data)
        reflection = self.llm(prompt)
        return {"reflection": reflection.strip()}


class RewritingNode(BaseNode):
    """
    Node that rewrites the input text according to a specified style or instruction.
    """

    def __init__(self, node_id: str, style: str = "formal"):
        super().__init__(node_id)
        self.style = style
        self.llm = get_llm()

    def process(self, input_data: Dict[str, str]) -> Dict[str, str]:
        # Expecting input_data to contain 'reflection' key
        reflection = input_data.get("reflection", "")
        prompt = (
            f"Rewrite the following reflection in a {self.style} style:\n\n{reflection}\n\nRewritten:"
        )
        rewritten = self.llm(prompt)
        return {"rewritten": rewritten.strip()}