# Copyright 2026 Flyto2. Licensed under Apache-2.0. See LICENSE.

"""
Local Ollama Integration Module

Provides local LLM support via Ollama for completely offline AI agent execution.
"""
import logging
from typing import Any, Dict

from ...base import BaseModule
from ...registry import register_module
from ....constants import OLLAMA_DEFAULT_URL


logger = logging.getLogger(__name__)


@register_module(
    module_id='ai.local_ollama.chat',
    version='1.0.0',
    category='ai',
    subcategory='ai',
    tags=['ai', 'ollama', 'local', 'llm', 'offline', 'chat', 'network', 'ssrf_protected'],
    label='Local Ollama Chat',
    label_key='modules.ai.local_ollama.chat.label',
    description='Chat with local LLM via Ollama (completely offline)',
    description_key='modules.ai.local_ollama.chat.description',
    icon='Server',
    color='#000000',

    # Connection types
    input_types=['text', 'json'],
    output_types=['text', 'json'],
    can_connect_to=['data.*', 'notify.*', 'file.*'],
    can_receive_from=['data.*', 'string.*', 'file.*', 'http.*', 'flow.*'],

    # Phase 2: Execution settings
    timeout_ms=120000,  # Local LLM can take time depending on hardware
    retryable=True,
    max_retries=3,
    concurrent_safe=False,  # Local LLM may have resource contention

    # Phase 2: Security settings
    requires_credentials=False,  # No API key needed for local
    handles_sensitive_data=True,  # User prompts may contain PII
    required_permissions=['shell.execute'],  # For localhost connection

    params_schema={
        'prompt': {
            'type': 'string',
            'label': 'Prompt',
            'label_key': 'modules.ai.local_ollama.chat.params.prompt.label',
            'description': 'The message to send to the local LLM',
            'description_key': 'modules.ai.local_ollama.chat.params.prompt.description',
            'required': True,
            'multiline': True
        ,
            'placeholder': 'Enter your prompt...',
},
        'model': {
            'type': 'select',
            'label': 'Model',
            'label_key': 'modules.ai.local_ollama.chat.params.model.label',
            'description': 'Ollama model to use',
            'description_key': 'modules.ai.local_ollama.chat.params.model.description',
            'options': [
                {'label': 'Llama 2 7B', 'value': 'llama2'},
                {'label': 'Llama 2 13B', 'value': 'llama2:13b'},
                {'label': 'Llama 2 70B', 'value': 'llama2:70b'},
                {'label': 'Mistral 7B', 'value': 'mistral'},
                {'label': 'Mixtral 8x7B', 'value': 'mixtral'},
                {'label': 'CodeLlama 7B', 'value': 'codellama'},
                {'label': 'CodeLlama 13B', 'value': 'codellama:13b'},
                {'label': 'Phi-2', 'value': 'phi'},
                {'label': 'Neural Chat 7B', 'value': 'neural-chat'},
                {'label': 'Starling 7B', 'value': 'starling-lm'}
            ],
            'default': 'llama2',
            'required': False
        },
        'temperature': {
            'type': 'number',
            'label': 'Temperature',
            'label_key': 'modules.ai.local_ollama.chat.params.temperature.label',
            'description': 'Sampling temperature (0-2)',
            'description_key': 'modules.ai.local_ollama.chat.params.temperature.description',
            'default': 0.7,
            'min': 0,
            'max': 2,
            'required': False
        },
        'system_message': {
            'type': 'string',
            'label': 'System Message',
            'label_key': 'modules.ai.local_ollama.chat.params.system_message.label',
            'description': 'System role message (optional)',
            'description_key': 'modules.ai.local_ollama.chat.params.system_message.description',
            'required': False,
            'multiline': True
        ,
            'placeholder': 'You are a helpful assistant.',
},
        'ollama_url': {
            'type': 'string',
            'label': 'Ollama URL',
            'label_key': 'modules.ai.local_ollama.chat.params.ollama_url.label',
            'description': 'Ollama server URL',
            'description_key': 'modules.ai.local_ollama.chat.params.ollama_url.description',
            'default': OLLAMA_DEFAULT_URL,
            'required': False
        ,
            'placeholder': 'http://localhost:11434',
},
        'max_tokens': {
            'type': 'number',
            'label': 'Max Tokens',
            'label_key': 'modules.ai.local_ollama.chat.params.max_tokens.label',
            'description': 'Maximum tokens in response (optional, depends on model)',
            'description_key': 'modules.ai.local_ollama.chat.params.max_tokens.description',
            'required': False
        }
    },
    output_schema={
        'response': {'type': 'string', 'description': 'Response from the operation',
                'description_key': 'modules.ai.local_ollama.chat.output.response.description'},
        'model': {'type': 'string', 'description': 'Model name or identifier',
                'description_key': 'modules.ai.local_ollama.chat.output.model.description'},
        'context': {'type': 'array', 'description': 'Conversation context for follow-up requests',
                'description_key': 'modules.ai.local_ollama.chat.output.context.description', 'items': {'type': 'number'}},
        'total_duration': {'type': 'number', 'description': 'Total processing duration',
                'description_key': 'modules.ai.local_ollama.chat.output.total_duration.description'},
        'load_duration': {'type': 'number', 'description': 'Model loading duration',
                'description_key': 'modules.ai.local_ollama.chat.output.load_duration.description'},
        'prompt_eval_count': {'type': 'number', 'description': 'Number of prompt tokens evaluated',
                'description_key': 'modules.ai.local_ollama.chat.output.prompt_eval_count.description'},
        'eval_count': {'type': 'number', 'description': 'Number of tokens generated',
                'description_key': 'modules.ai.local_ollama.chat.output.eval_count.description'}
    },
    examples=[
        {
            'title': 'Simple local chat',
            'params': {
                'prompt': 'Explain quantum computing in 3 sentences',
                'model': 'llama2'
            }
        },
        {
            'title': 'Code generation with local model',
            'params': {
                'prompt': 'Write a Python function to calculate fibonacci numbers',
                'model': 'codellama',
                'temperature': 0.2,
                'system_message': 'You are a Python programming expert. Write clean, efficient code.'
            }
        },
        {
            'title': 'Local reasoning task',
            'params': {
                'prompt': 'What are the pros and cons of microservices architecture?',
                'model': 'mistral',
                'temperature': 0.7
            }
        }
    ],
    author='Flyto Team',
    license='MIT'
)
class LocalOllamaChatModule(BaseModule):
    """Local Ollama Chat Module - Completely offline LLM"""

    def validate_params(self) -> None:
        import os
        from urllib.parse import urlparse

        self.prompt = self.params.get('prompt')
        self.model = self.params.get('model', 'llama2')
        self.temperature = self.params.get('temperature', 0.7)
        self.system_message = self.params.get('system_message')
        self.ollama_url = self.params.get('ollama_url', OLLAMA_DEFAULT_URL)
        self.max_tokens = self.params.get('max_tokens')

        if not self.prompt:
            raise ValueError("prompt is required")

        # SECURITY: Validate ollama_url - restrict to localhost by default
        # This module is for LOCAL Ollama, not remote servers
        parsed = urlparse(self.ollama_url)
        host = parsed.hostname or ''
        localhost_hosts = ('localhost', '127.0.0.1', '::1', '0.0.0.0')

        # Allow localhost connections always
        if host.lower() not in localhost_hosts:
            # Check if remote Ollama is explicitly allowed
            allow_remote = os.environ.get('FLYTO_ALLOW_REMOTE_OLLAMA', '').lower() == 'true'
            if not allow_remote:
                raise ValueError(
                    f"SSRF blocked: ollama_url must be localhost (got {host}). "
                    "Set FLYTO_ALLOW_REMOTE_OLLAMA=true to allow remote servers."
                )

    async def execute(self) -> Any:
        try:
            import aiohttp
            import json

            # Build messages
            messages = []
            if self.system_message:
                messages.append({
                    "role": "system",
                    "content": self.system_message
                })
            messages.append({
                "role": "user",
                "content": self.prompt
            })

            # Build request payload
            payload = {
                "model": self.model,
                "messages": messages,
                "stream": False,
                "options": {
                    "temperature": self.temperature
                }
            }

            # Add max_tokens if specified
            if self.max_tokens:
                payload["options"]["num_predict"] = self.max_tokens

            # Make API call to local Ollama
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.ollama_url}/api/chat",
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=120)
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        raise RuntimeError(
                            f"Ollama API error (status {response.status}): {error_text}"
                        )

                    result = await response.json()

            # Extract response
            message = result.get('message', {})
            content = message.get('content', '')

            return {
                "response": content,
                "model": result.get('model', self.model),
                "context": result.get('context', []),
                "total_duration": result.get('total_duration', 0),
                "load_duration": result.get('load_duration', 0),
                "prompt_eval_count": result.get('prompt_eval_count', 0),
                "eval_count": result.get('eval_count', 0)
            }

        except aiohttp.ClientError as e:
            raise RuntimeError(
                f"Failed to connect to Ollama at {self.ollama_url}. "
                f"Make sure Ollama is running (ollama serve). Error: {str(e)}"
            )
        except Exception as e:
            raise RuntimeError(f"Local Ollama error: {str(e)}")
