import logging
import os
from typing import Any

import json_repair
from jinja2 import Template

from app.core.error_codes import BizCode
from app.core.exceptions import BusinessException
from app.core.models import RedBearLLM, RedBearModelConfig
from app.core.workflow.engine.state_manager import WorkflowState
from app.core.workflow.engine.variable_pool import VariablePool
from app.core.workflow.nodes.base_node import BaseNode
from app.core.workflow.nodes.parameter_extractor.config import ParameterExtractorNodeConfig
from app.core.workflow.variable.base_variable import VariableType
from app.db import get_db_read
from app.models import ModelType
from app.services.model_service import ModelConfigService

logger = logging.getLogger(__name__)


class ParameterExtractorNode(BaseNode):
    def __init__(self, node_config: dict[str, Any], workflow_config: dict[str, Any]):
        super().__init__(node_config, workflow_config)
        self.typed_config: ParameterExtractorNodeConfig | None = None
        self.response_metadata = {}

    def _extract_token_usage(self, business_result: Any) -> dict[str, int] | None:
        if self.response_metadata:
            usage = self.response_metadata.get('token_usage')
            if usage:
                return {
                    "prompt_tokens": usage.get('input_tokens', 0),
                    "completion_tokens": usage.get('output_tokens', 0),
                    "total_tokens": usage.get('total_tokens', 0)
                }
        return None

    def _extract_input(self, state: WorkflowState, variable_pool: VariablePool) -> dict[str, Any]:
        return {
            "text": self._render_template(self.typed_config.text, variable_pool),
            "prompt": self._render_template(self.typed_config.prompt, variable_pool),
            "params": [param.model_dump(mode="json") for param in self.typed_config.params],
            "model_id": str(self.typed_config.model_id),
        }

    def _output_types(self) -> dict[str, VariableType]:
        outputs = {}
        for param in self.typed_config.params:
            outputs[param.name] = param.type
        return outputs

    @staticmethod
    def _get_prompt():
        """
        Load system and user prompt templates from local prompt files.

        Notes:
        - Templates are expected to be Jinja2 files.
        - Reading from disk each time ensures the latest template is used (could be cached if performance-critical).
        - Both templates must exist, otherwise an exception will be raised.

        Returns:
            Tuple[str, str]: system_prompt, user_prompt
        """
        current_dir = os.path.dirname(os.path.abspath(__file__))
        with open(
                os.path.join(current_dir, "prompt", "system_prompt.jinja2"),
                encoding='utf-8'
        ) as f:
            system_prompt = f.read()
        with open(os.path.join(
                current_dir, "prompt", "user_prompt.jinja2"),
                encoding='utf-8'
        ) as f:
            user_prompt = f.read()
        return system_prompt, user_prompt

    def _get_llm_instance(self) -> RedBearLLM:
        """
        Retrieve a configured LLM instance based on the model ID from database.

        Responsibilities:
        - Validate that the model exists and has at least one API key configured.
        - Construct RedBearLLM instance with proper credentials and model type.
        - Raise clear BusinessException if configuration is invalid.

        Returns:
            RedBearLLM: Configured LLM instance ready to be invoked.

        Raises:
            BusinessException: If the model is missing or lacks valid API key.
        """
        model_id = self.typed_config.model_id

        with get_db_read() as db:
            config = ModelConfigService.get_model_by_id(db=db, model_id=model_id)

            if not config:
                raise BusinessException("Configured model does not exist", BizCode.NOT_FOUND)

            if not config.api_keys or len(config.api_keys) == 0:
                raise BusinessException("Model configuration is missing API Key", BizCode.INVALID_PARAMETER)

            api_config = self.model_balance(config)
            model_name = api_config.model_name
            provider = api_config.provider
            api_key = api_config.api_key
            api_base = api_config.api_base
            is_omni = api_config.is_omni
            model_type = config.type

        llm = RedBearLLM(
            RedBearModelConfig(
                model_name=model_name,
                provider=provider,
                api_key=api_key,
                base_url=api_base,
                is_omni=is_omni
            ),
            type=ModelType(model_type)
        )
        return llm

    def _get_field_desc(self) -> dict[str, str]:
        """
        Build a dictionary mapping each parameter name to its description.
        Useful for dynamically generating prompts for LLM.

        Returns:
            dict[str, str]: Mapping of parameter names to descriptions.
        """
        field_desc = {}
        for param in self.typed_config.params:
            field_desc[param.name] = param.desc
        return field_desc

    def _get_field_type(self) -> dict[str, str]:
        """
        Build a dictionary mapping each parameter name to its description.
        Useful for dynamically generating prompts for LLM.

        Returns:
            dict[str, str]: Mapping of parameter names to descriptions.
        """
        field_type = {}
        for param in self.typed_config.params:
            field_type[param.name] = f'{param.type}, required:{str(param.required)}'
        return field_type

    async def execute(self, state: WorkflowState, variable_pool: VariablePool) -> Any:
        """
        Main execution function for this node.

        Workflow:
        1. Retrieve LLM instance with valid credentials.
        2. Render user prompt template with field descriptions, types, and input text.
        3. Send system and user prompts to LLM asynchronously.
        4. Repair LLM JSON output safely.
        5. Return output dictionary.

        Notes:
        - JSON repair is used to handle minor formatting errors in LLM output.
        - Exceptions are raised explicitly if parsing fails, to prevent silent workflow failures.
        - Rendering uses self._render_template for dynamic substitution from workflow state.

        Args:
            state (WorkflowState): Current state of the workflow, used for template rendering.
            variable_pool (VariablePool): Used for accessing and setting variables during execution.

        Returns:
            dict[str, Any]: Dictionary containing extracted parameters under the "output" key.

        Raises:
            BusinessException: If LLM output cannot be parsed as valid JSON.
        """
        self.typed_config = ParameterExtractorNodeConfig(**self.config)
        llm = self._get_llm_instance()
        system_prompt, user_prompt = self._get_prompt()

        user_prompt_teplate = Template(user_prompt)
        rendered_user_prompt = user_prompt_teplate.render(
            field_descriptions=str(self._get_field_desc()),
            field_type=str(self._get_field_type()),
            text_input=self._render_template(self.typed_config.text, variable_pool)
        )

        messages = [
            ("system", system_prompt),

        ]
        if self.typed_config.prompt:
            messages.extend([
                ("user", self._render_template(self.typed_config.prompt, variable_pool)),
                ("user", rendered_user_prompt),
            ])
        else:
            messages.extend([
                ("user", rendered_user_prompt),
            ])

        model_resp = await llm.ainvoke(messages)
        self.response_metadata = model_resp.response_metadata
        model_message = self.process_model_output(model_resp.content)
        result = json_repair.repair_json(model_message, return_objects=True)
        logger.info(f"node: {self.node_id} get params:{result}")

        return result
