#!/usr/bin/env python3
"""
LLM Client
==========

Client for interacting with LLM providers (Ollama, Claude API).
Handles request formatting, retry logic, timeout handling, and fallback.

Supported Providers:
- Ollama (local): Default, uses llama3:8b model
- Claude API: Fallback when Ollama fails

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import os
import time
import json
from typing import Optional, Dict, Any, Callable
from dataclasses import dataclass
from abc import ABC, abstractmethod
from loguru import logger

# HTTP client - try httpx first, fall back to requests
try:
    import httpx
    HAS_HTTPX = True
except ImportError:
    HAS_HTTPX = False

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

if not HAS_HTTPX and not HAS_REQUESTS:
    logger.error("Neither httpx nor requests available - LLM client will not work")


# ============================================================
# CONFIGURATION
# ============================================================

@dataclass
class LLMClientConfig:
    """Configuration for LLM client"""
    # Ollama settings
    ollama_host: str = "http://localhost:11434"
    ollama_model: str = "llama3:8b"
    ollama_timeout: int = 120  # seconds

    # Claude API settings
    claude_api_key: Optional[str] = None
    claude_model: str = "claude-3-haiku-20240307"
    claude_timeout: int = 60  # seconds

    # Retry settings
    max_retries: int = 3
    retry_delay: float = 1.0  # seconds
    retry_backoff: float = 2.0  # multiplier

    # Fallback
    enable_fallback: bool = True


# ============================================================
# ABSTRACT BASE CLASS
# ============================================================

class LLMProvider(ABC):
    """Abstract base class for LLM providers"""

    @abstractmethod
    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7
    ) -> str:
        """Generate a response from the LLM"""
        pass

    @abstractmethod
    def is_available(self) -> bool:
        """Check if the provider is available"""
        pass


# ============================================================
# OLLAMA PROVIDER
# ============================================================

class OllamaProvider(LLMProvider):
    """
    Ollama LLM provider for local inference.

    Uses the Ollama API to run local LLMs like Llama 3.

    Requirements:
    - Ollama running locally (ollama serve)
    - Model pulled (ollama pull llama3:8b)
    """

    def __init__(
        self,
        host: str = "http://localhost:11434",
        model: str = "llama3:8b",
        timeout: int = 120
    ):
        self.host = host.rstrip("/")
        self.model = model
        self.timeout = timeout
        self._available = None  # Cache availability check

    def is_available(self) -> bool:
        """Check if Ollama is running and model is available"""
        if self._available is not None:
            return self._available

        try:
            url = f"{self.host}/api/tags"

            if HAS_HTTPX:
                with httpx.Client(timeout=5) as client:
                    response = client.get(url)
            elif HAS_REQUESTS:
                response = requests.get(url, timeout=5)
            else:
                return False

            if response.status_code == 200:
                data = response.json()
                models = [m.get("name", "") for m in data.get("models", [])]
                # Check if our model (or a variant) is available
                model_base = self.model.split(":")[0]
                self._available = any(model_base in m for m in models)
                if self._available:
                    logger.info(f"Ollama available with model {self.model}")
                else:
                    logger.warning(f"Ollama running but model {self.model} not found")
                return self._available

        except Exception as e:
            logger.debug(f"Ollama not available: {e}")

        self._available = False
        return False

    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7
    ) -> str:
        """
        Generate a response using Ollama.

        Args:
            prompt: User prompt
            system_prompt: System prompt (optional)
            temperature: Sampling temperature

        Returns:
            Generated response text

        Raises:
            RuntimeError: If generation fails
        """
        url = f"{self.host}/api/generate"

        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": 1024,  # Max tokens
            }
        }

        if system_prompt:
            payload["system"] = system_prompt

        try:
            if HAS_HTTPX:
                with httpx.Client(timeout=self.timeout) as client:
                    response = client.post(url, json=payload)
            elif HAS_REQUESTS:
                response = requests.post(url, json=payload, timeout=self.timeout)
            else:
                raise RuntimeError("No HTTP client available")

            if response.status_code != 200:
                raise RuntimeError(f"Ollama error: {response.status_code} - {response.text}")

            data = response.json()
            return data.get("response", "")

        except Exception as e:
            raise RuntimeError(f"Ollama generation failed: {e}")


# ============================================================
# CLAUDE PROVIDER
# ============================================================

class ClaudeProvider(LLMProvider):
    """
    Claude API provider for cloud inference.

    Uses the Anthropic Claude API.

    Requirements:
    - ANTHROPIC_API_KEY environment variable
    """

    API_URL = "https://api.anthropic.com/v1/messages"

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "claude-3-haiku-20240307",
        timeout: int = 60
    ):
        self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
        self.model = model
        self.timeout = timeout

    def is_available(self) -> bool:
        """Check if Claude API key is configured"""
        return bool(self.api_key)

    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7
    ) -> str:
        """
        Generate a response using Claude API.

        Args:
            prompt: User prompt
            system_prompt: System prompt (optional)
            temperature: Sampling temperature

        Returns:
            Generated response text

        Raises:
            RuntimeError: If generation fails
        """
        if not self.api_key:
            raise RuntimeError("Claude API key not configured")

        headers = {
            "Content-Type": "application/json",
            "x-api-key": self.api_key,
            "anthropic-version": "2023-06-01"
        }

        payload = {
            "model": self.model,
            "max_tokens": 1024,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature
        }

        if system_prompt:
            payload["system"] = system_prompt

        try:
            if HAS_HTTPX:
                with httpx.Client(timeout=self.timeout) as client:
                    response = client.post(self.API_URL, headers=headers, json=payload)
            elif HAS_REQUESTS:
                response = requests.post(
                    self.API_URL,
                    headers=headers,
                    json=payload,
                    timeout=self.timeout
                )
            else:
                raise RuntimeError("No HTTP client available")

            if response.status_code != 200:
                raise RuntimeError(f"Claude API error: {response.status_code} - {response.text}")

            data = response.json()
            content = data.get("content", [])

            if content and len(content) > 0:
                return content[0].get("text", "")

            return ""

        except Exception as e:
            raise RuntimeError(f"Claude generation failed: {e}")


# ============================================================
# LLM CLIENT (MAIN CLASS)
# ============================================================

class LLMClient:
    """
    Main LLM client with provider abstraction and fallback.

    Usage:
        client = LLMClient(provider="ollama", model="llama3:8b")

        response = client.get_trading_decision(prompt)

        # Or with automatic fallback
        client = LLMClient(enable_fallback=True)
        response = client.get_trading_decision(prompt)
    """

    def __init__(
        self,
        provider: str = "ollama",
        model: Optional[str] = None,
        config: Optional[LLMClientConfig] = None,
        enable_fallback: bool = True
    ):
        """
        Initialize the LLM client.

        Args:
            provider: Primary provider ("ollama" or "claude")
            model: Model to use (provider-specific default if None)
            config: Full configuration object (optional)
            enable_fallback: Enable fallback to other providers
        """
        self.config = config or LLMClientConfig()
        self.enable_fallback = enable_fallback

        # Initialize providers
        self.providers: Dict[str, LLMProvider] = {}

        # Set up Ollama
        ollama_model = model if provider == "ollama" and model else self.config.ollama_model
        self.providers["ollama"] = OllamaProvider(
            host=self.config.ollama_host,
            model=ollama_model,
            timeout=self.config.ollama_timeout
        )

        # Set up Claude
        claude_model = model if provider == "claude" and model else self.config.claude_model
        self.providers["claude"] = ClaudeProvider(
            api_key=self.config.claude_api_key,
            model=claude_model,
            timeout=self.config.claude_timeout
        )

        # Set primary provider
        self.primary_provider = provider

        # Determine fallback order
        if provider == "ollama":
            self.fallback_order = ["ollama", "claude"]
        else:
            self.fallback_order = ["claude", "ollama"]

        logger.info(f"LLM client initialized with primary provider: {provider}")

    def _get_available_provider(self) -> Optional[LLMProvider]:
        """Get the first available provider"""
        for provider_name in self.fallback_order:
            provider = self.providers.get(provider_name)
            if provider and provider.is_available():
                return provider

        return None

    def _generate_with_retry(
        self,
        provider: LLMProvider,
        prompt: str,
        system_prompt: Optional[str],
        temperature: float
    ) -> str:
        """Generate with retry logic"""
        last_error = None
        delay = self.config.retry_delay

        for attempt in range(self.config.max_retries):
            try:
                return provider.generate(prompt, system_prompt, temperature)
            except Exception as e:
                last_error = e
                logger.warning(f"Attempt {attempt + 1} failed: {e}")

                if attempt < self.config.max_retries - 1:
                    time.sleep(delay)
                    delay *= self.config.retry_backoff

        raise RuntimeError(f"All retry attempts failed: {last_error}")

    def get_trading_decision(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.3  # Lower for more consistent trading decisions
    ) -> str:
        """
        Get a trading decision from the LLM.

        Args:
            prompt: The trading analysis prompt
            system_prompt: System prompt (optional, uses default if None)
            temperature: Sampling temperature (default 0.3 for consistency)

        Returns:
            Raw LLM response text

        Raises:
            RuntimeError: If all providers fail
        """
        # Import default system prompt if needed
        if system_prompt is None:
            from .prompts.trading_decision import SYSTEM_PROMPT
            system_prompt = SYSTEM_PROMPT

        providers_to_try = self.fallback_order if self.enable_fallback else [self.primary_provider]

        for provider_name in providers_to_try:
            provider = self.providers.get(provider_name)

            if not provider or not provider.is_available():
                logger.debug(f"Provider {provider_name} not available, skipping")
                continue

            try:
                logger.info(f"Generating with {provider_name}...")
                response = self._generate_with_retry(
                    provider, prompt, system_prompt, temperature
                )

                if response:
                    logger.debug(f"Got response from {provider_name}")
                    return response

            except Exception as e:
                logger.warning(f"Provider {provider_name} failed: {e}")

                if not self.enable_fallback:
                    raise

        raise RuntimeError("All LLM providers failed")

    def check_availability(self) -> Dict[str, bool]:
        """
        Check availability of all providers.

        Returns:
            Dict mapping provider names to availability status
        """
        return {
            name: provider.is_available()
            for name, provider in self.providers.items()
        }

    def test_connection(self) -> bool:
        """
        Test connection to the primary provider.

        Returns:
            True if connection successful
        """
        provider = self.providers.get(self.primary_provider)

        if not provider:
            return False

        try:
            if not provider.is_available():
                return False

            # Try a simple generation
            response = provider.generate(
                "Say 'OK' if you're working.",
                temperature=0.0
            )

            return bool(response)

        except Exception as e:
            logger.error(f"Connection test failed: {e}")
            return False


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

def create_client(
    provider: str = "ollama",
    model: Optional[str] = None
) -> LLMClient:
    """
    Create an LLM client with default settings.

    Args:
        provider: Provider name ("ollama" or "claude")
        model: Model name (optional)

    Returns:
        Configured LLMClient
    """
    return LLMClient(provider=provider, model=model)


def get_ollama_client(model: str = "llama3:8b") -> LLMClient:
    """Create an Ollama client"""
    return LLMClient(provider="ollama", model=model)


def get_claude_client(model: str = "claude-3-haiku-20240307") -> LLMClient:
    """Create a Claude client"""
    return LLMClient(provider="claude", model=model)


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    'LLMClient',
    'LLMClientConfig',
    'LLMProvider',
    'OllamaProvider',
    'ClaudeProvider',
    'create_client',
    'get_ollama_client',
    'get_claude_client',
]
