"""Base LLM provider interface.

This module defines the base interface for LLM providers. Implementations
should extend BaseLLMProvider and implement the abstract methods.

The interface is designed around JSON generation, which is the primary
use case for company resolution and normalization tasks.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)


class LLMError(Exception):
    """Base exception for LLM provider errors."""


@dataclass
class LLMResponse:
    """Response from an LLM generation request.

    Attributes:
        value: The generated content (usually parsed JSON dict)
        raw: Raw response text before parsing
        model: Model identifier used
        tokens_in: Input token count (if available)
        tokens_out: Output token count (if available)
        latency_ms: Request latency in milliseconds
        metrics: Additional provider-specific metrics
    """
    value: Dict[str, Any]
    raw: str = ""
    model: str = ""
    tokens_in: int = 0
    tokens_out: int = 0
    latency_ms: float = 0.0
    metrics: Dict[str, Any] = field(default_factory=dict)


class BaseLLMProvider(ABC):
    """Abstract base class for LLM providers.

    Implementations should handle:
    - Authentication/credentials for the specific provider
    - Request formatting for the provider's API
    - Response parsing and error handling
    - JSON schema enforcement (if supported by provider)

    The primary method is `generate_json()` which generates structured
    JSON output from a prompt.

    Example implementation:
        class VertexAIProvider(BaseLLMProvider):
            def generate_json(self, prompt, schema=None, **options):
                # Call Vertex AI API
                response = self._call_vertex_api(prompt, schema)
                return {"value": response.parsed_json, ...}
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Provider name (e.g., 'vertex_ai', 'ollama')."""
        ...

    @property
    def model(self) -> str:
        """Current model identifier."""
        return ""

    @abstractmethod
    def generate_json(
        self,
        prompt: str,
        schema: Optional[Dict[str, Any]] = None,
        **options: Any,
    ) -> Dict[str, Any]:
        """Generate structured JSON from a prompt.

        This is the primary generation method. It should return a dict
        with at least a 'value' key containing the parsed JSON.

        Args:
            prompt: The prompt to send to the LLM
            schema: Optional JSON schema for response validation.
                   Providers that support structured output should use
                   this to constrain generation.
            **options: Provider-specific options like:
                - temperature: Sampling temperature (0.0-1.0)
                - max_tokens: Maximum tokens to generate
                - num_predict: Alias for max_tokens

        Returns:
            Dict with keys:
            - value: Parsed JSON response (dict)
            - metrics: Optional generation metrics

        Raises:
            LLMError: If generation fails
        """
        ...

    def generate_text(
        self,
        prompt: str,
        **options: Any,
    ) -> str:
        """Generate plain text from a prompt.

        Default implementation calls generate_json and extracts text.
        Override for providers with native text generation.

        Args:
            prompt: The prompt to send to the LLM
            **options: Provider-specific options

        Returns:
            Generated text string
        """
        result = self.generate_json(prompt, **options)
        value = result.get("value", {})
        if isinstance(value, dict):
            # Try common text keys
            for key in ("text", "content", "response", "output"):
                if key in value:
                    return str(value[key])
            # Fall back to first string value
            for v in value.values():
                if isinstance(v, str):
                    return v
        return str(value)

    def is_available(self) -> bool:
        """Check if the provider is available and configured.

        Default implementation returns True. Override to add
        provider-specific availability checks.
        """
        return True

    def close(self) -> None:
        """Release any resources held by the provider.

        Default implementation does nothing. Override if the provider
        maintains connections or other resources.
        """
        pass

    def __enter__(self) -> "BaseLLMProvider":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()
