from __future__ import annotations
import traceback
from loguru import logger

from service_forge.workflow.node import Node
from service_forge.workflow.port import Port
from service_forge.workflow.registry.prompt_registry import load_prompt
from service_forge.llm import Model
from service_forge.llm.cleaner.json_cleaner import JsonCleaner

class LLMNode(Node):
    DEFAULT_INPUT_PORTS = [
        Port("prompt", str),
        Port("system_prompt", str),
        Port("model", Model),
        Port("temperature", float),
        Port("cleaner", str),
    ]

    DEFAULT_OUTPUT_PORTS = [
        Port("content", str),
        Port("error", str),
    ]

    def __init__(
        self, 
        name: str,
    ) -> None:
        super().__init__(
            name,
        )

    async def _run(
        self,
        prompt: str,
        system_prompt: str,
        model: Model,
        temperature: float,
        cleaner: str,
    ) -> None:
        try:
            system_prompt = load_prompt(system_prompt)
            prompt = load_prompt(prompt)

            response = await self.llm.acompletion(
                input=prompt,
                system_prompt=system_prompt,
                model=model,
                temperature=temperature,
            )

            content = response.choices[0].message.content

            if cleaner:
                for name in cleaner.split('|'):
                    if name.strip().lstrip() == 'json':
                        content = JsonCleaner().clean(content)
        
            self.activate_output_edges('content', content)
        except Exception as e:
            logger.error(f"Error when running LLMNode {self.name}: {e}, traceback: {traceback.format_exc()}")
            self.activate_output_edges('error', str(e))