"""
LLM integration nodes for AI model interactions.
"""

from typing import Any, Dict, Optional
from ybagent.node import BaseNode, NodeType, NodeInput, NodeOutput


class LLMPromptNode(BaseNode):
    """
    Node for managing and formatting prompts for LLMs.
    """
    
    def __init__(self, id: str, template: str = "", **config):
        """
        Initialize prompt node.
        
        Args:
            id: Node ID
            template: Prompt template with {variable} placeholders
            **config: Additional configuration
        """
        config["template"] = template
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.LLM, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="variables", type="dict", required=False, description="Variables for template")
        ]
        self.metadata.outputs = [
            NodeOutput(name="prompt", type="str", description="Formatted prompt")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Format prompt template with variables.
        
        Args:
            inputs: Input variables
            context: Execution context
        
        Returns:
            Formatted prompt
        """
        template = self.metadata.config.get("template", "")
        variables = inputs.get("variables", {})
        
        try:
            prompt = template.format(**variables)
        except KeyError as e:
            raise ValueError(f"Missing variable in template: {e}")
        
        return {"prompt": prompt}


class OllamaNode(BaseNode):
    """
    Node for Ollama local LLM integration.
    
    Ollama allows you to run LLMs locally without API keys.
    Install Ollama from: https://ollama.ai
    """
    
    def __init__(
        self,
        id: str,
        model: str = "llama2",
        base_url: str = "http://localhost:11434",
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **config
    ):
        """
        Initialize Ollama node.
        
        Args:
            id: Node ID
            model: Model name (e.g., 'llama2', 'mistral', 'codellama')
            base_url: Ollama server URL
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            **config: Additional configuration
        """
        config.update({
            "model": model,
            "base_url": base_url,
            "temperature": temperature,
            "max_tokens": max_tokens
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.LLM, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="prompt", type="str", required=True, description="Prompt for the LLM"),
            NodeInput(name="system_message", type="str", required=False, description="System message"),
            NodeInput(name="context", type="any", required=False, description="Retrieval context")
        ]
        self.metadata.outputs = [
            NodeOutput(name="response", type="str", description="LLM response"),
            NodeOutput(name="output", type="str", description="LLM response")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call Ollama API.
        
        Args:
            inputs: Prompt and optional system message
            context: Execution context
        
        Returns:
            LLM response
        """
        try:
            import httpx
        except ImportError:
            raise ImportError("httpx package not installed. Install with: pip install httpx")
        
        model = self.metadata.config.get("model", "llama2")
        base_url = self.metadata.config.get("base_url", "http://localhost:11434")
        temperature = self.metadata.config.get("temperature", 0.7)
        max_tokens = self.metadata.config.get("max_tokens")
        
        # Build messages
        messages = []
        
        system_message_input = inputs.get("system_message")
        system_prompt_cfg = self.metadata.config.get("system_prompt")
        system_message = system_message_input or system_prompt_cfg
        if system_message:
            messages.append({"role": "system", "content": system_message})
        
        # Add user prompt
        prompt = inputs.get("prompt", "")
        original_query = prompt  # Store original query for display
        if isinstance(prompt, dict):
            if "query_input" in prompt and isinstance(prompt["query_input"], str):
                original_query = prompt["query_input"]
                prompt = prompt["query_input"]
            elif "prompt" in prompt and isinstance(prompt["prompt"], str):
                original_query = prompt["prompt"]
                prompt = prompt["prompt"]
            else:
                try:
                    prompt = next((str(v) for v in prompt.values() if isinstance(v, str)), str(prompt))
                    original_query = prompt
                except Exception:
                    prompt = str(prompt)
                    original_query = prompt
        elif isinstance(prompt, list):
            prompt = "\n".join(str(x) for x in prompt)
            original_query = prompt
        ctx = inputs.get("context")
        if ctx is not None:
            if isinstance(ctx, list):
                ctx_text = "\n".join([str(x) for x in ctx])
            elif isinstance(ctx, dict):
                ctx_text = "\n".join([f"{k}: {v}" for k, v in ctx.items()])
            else:
                ctx_text = str(ctx)
            prompt = f"Context:\n{ctx_text}\n\nQuestion:\n{prompt}"
        messages.append({"role": "user", "content": prompt})
        
        # Build options
        options = {"temperature": temperature}
        if max_tokens:
            options["num_predict"] = max_tokens
        
        # Call Ollama API
        try:
            with httpx.Client(timeout=120.0) as client:
                response = client.post(
                    f"{base_url}/api/chat",
                    json={
                        "model": model,
                        "messages": messages,
                        "stream": False,
                        "options": options
                    }
                )
            
            if response.status_code != 200:
                raise RuntimeError(
                    f"Ollama API error: {response.status_code} - {response.text}\n"
                    f"Make sure Ollama is running (ollama serve) and the model '{model}' is installed (ollama pull {model})"
                )
            
            result = response.json()
            text = result["message"]["content"]

            try:
                from ybagent.utils.display import DisplayManager
                state = context.get("state") if isinstance(context, dict) else None
                suppress_llm = state.get_variable("suppress_llm_response", False) if state else False
                suppress_vector = state.get_variable("suppress_vector_display", False) if state else False
                if self.metadata.config.get("display_response", True) and not suppress_llm:
                    DisplayManager.llm_response(text, title="AI Response", query=original_query)
            except ImportError:
                pass
            
            return {"response": text, "output": text}
            
        except httpx.ConnectError:
            raise RuntimeError(
                f"Cannot connect to Ollama at {base_url}. "
                "Make sure Ollama is running with 'ollama serve'"
            )


class OpenAINode(BaseNode):
    """
    Node for OpenAI API integration.
    """
    
    def __init__(
        self,
        id: str,
        api_key: str = "",
        model: str = "gpt-4",
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **config
    ):
        """
        Initialize OpenAI node.
        
        Args:
            id: Node ID
            api_key: OpenAI API key
            model: Model name
            temperature: Sampling temperature
            max_tokens: Maximum tokens
            **config: Additional configuration
        """
        config.update({
            "api_key": api_key,
            "model": model,
            "temperature": temperature,
            "max_tokens": max_tokens
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.LLM, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="prompt", type="str", required=True, description="Prompt for the LLM"),
            NodeInput(name="system_message", type="str", required=False, description="System message")
        ]
        self.metadata.outputs = [
            NodeOutput(name="response", type="str", description="LLM response"),
            NodeOutput(name="usage", type="dict", description="Token usage information")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call OpenAI API.
        
        Args:
            inputs: Prompt and optional system message
            context: Execution context
        
        Returns:
            LLM response and usage information
        """
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError("OpenAI package not installed. Install with: pip install openai")
        
        api_key = self.metadata.config.get("api_key")
        if not api_key:
            # Try to get from env var
            import os
            api_key = os.environ.get("OPENAI_API_KEY")
        
        if not api_key:
            raise ValueError("OpenAI API key not provided")
        
        client = OpenAI(api_key=api_key)
        
        messages = []
        
        # Add system message if provided
        system_message = inputs.get("system_message")
        if system_message:
            messages.append({"role": "system", "content": system_message})
        
        messages.append({"role": "user", "content": inputs.get("prompt", "")})
        
        # Call API
        response = client.chat.completions.create(
            model=self.metadata.config.get("model", "gpt-4"),
            messages=messages,
            temperature=self.metadata.config.get("temperature", 0.7),
            max_tokens=self.metadata.config.get("max_tokens")
        )
        
        return {
            "response": response.choices[0].message.content,
            "usage": {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens
            }
        }


class HuggingFaceNode(BaseNode):
    """
    Node for HuggingFace integration (API or local).
    """
    
    def __init__(
        self,
        id: str,
        model_name: str = "gpt2",
        use_api: bool = True,
        api_key: str = "",
        **config
    ):
        """
        Initialize HuggingFace node.
        
        Args:
            id: Node ID
            model_name: Model identifier
            use_api: Whether to use Inference API or local transformers
            api_key: HuggingFace API key (for API usage)
            **config: Additional configuration
        """
        config.update({
            "model_name": model_name,
            "use_api": use_api,
            "api_key": api_key
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.LLM, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="prompt", type="str", required=True, description="Input prompt")
        ]
        self.metadata.outputs = [
            NodeOutput(name="response", type="str", description="Model response")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call HuggingFace model.
        
        Args:
            inputs: Input prompt
            context: Execution context
        
        Returns:
            Model response
        """
        prompt = inputs.get("prompt", "")
        model_name = self.metadata.config.get("model_name", "gpt2")
        use_api = self.metadata.config.get("use_api", True)
        
        if use_api:
            # Use HuggingFace Inference API
            try:
                import httpx
            except ImportError:
                raise ImportError("httpx package not installed. Install with: pip install httpx")
            
            api_key = self.metadata.config.get("api_key")
            if not api_key:
                raise ValueError("HuggingFace API key not provided")
            
            headers = {"Authorization": f"Bearer {api_key}"}
            api_url = f"https://api-inference.huggingface.co/models/{model_name}"
            
            response = httpx.post(
                api_url,
                headers=headers,
                json={"inputs": prompt},
                timeout=30.0
            )
            
            if response.status_code != 200:
                raise RuntimeError(f"HuggingFace API error: {response.text}")
            
            result = response.json()
            
            # Handle different response formats
            if isinstance(result, list) and len(result) > 0:
                if "generated_text" in result[0]:
                    return {"response": result[0]["generated_text"]}
                else:
                    return {"response": str(result[0])}
            else:
                return {"response": str(result)}
        
        else:
            # Use local transformers
            try:
                from transformers import pipeline
            except ImportError:
                raise ImportError(
                    "transformers package not installed. "
                    "Install with: pip install transformers torch"
                )
            
            generator = pipeline("text-generation", model=model_name)
            result = generator(prompt, max_length=100, num_return_sequences=1)
            
            return {"response": result[0]["generated_text"]}
