"""LLM provider for transformer fallback support.

Defaults to Vertex AI with Gemini 3 Flash. Ollama is deprecated.
"""

import logging
import os
import json
from typing import Dict, Any, Optional

logger = logging.getLogger(__name__)

# Try Vertex AI first (preferred)
try:
    from .providers.vertex_ai_provider import VertexAIProvider
    HAS_VERTEX = True
except ImportError:
    HAS_VERTEX = False
    VertexAIProvider = None  # type: ignore

# Optional Ollama dependency (deprecated)
try:
    import ollama
    HAS_OLLAMA = True
except ImportError:
    HAS_OLLAMA = False


class LLMProvider:
    """Base LLM provider interface."""

    def generate_json(
        self, prompt: str, schema: Optional[Dict] = None, num_predict: int = 256
    ) -> Dict[str, Any]:
        """Generate JSON response from prompt."""
        raise NotImplementedError


class VertexProvider(LLMProvider):
    """Vertex AI-based LLM provider using Gemini models."""

    def __init__(self, model: str = None):
        if not HAS_VERTEX:
            raise ImportError("Vertex AI SDK not installed")
        self.model = model or os.getenv("VERTEX_MODEL") or "gemini-3-flash-preview"
        self._provider = VertexAIProvider(model=self.model)

    def generate_json(
        self, prompt: str, schema: Optional[Dict] = None, num_predict: int = 256
    ) -> Dict[str, Any]:
        """Generate JSON using Vertex AI."""
        try:
            result = self._provider.generate_json(
                prompt,
                schema=schema,
                max_output_tokens=num_predict,
            )
            return {
                "value": result.get("value", {}),
                "metrics": {
                    "model": self.model,
                    "provider": "vertex_ai",
                    "latency_ms": result.get("metrics", {}).get("latency_ms", 0),
                },
            }
        except Exception as e:
            logger.warning("Vertex AI generation failed: %s", e)
            return {
                "value": {},
                "metrics": {"model": self.model, "provider": "vertex_ai", "error": str(e)},
            }


class OllamaProvider(LLMProvider):
    """Ollama-based LLM provider for local AI. DEPRECATED: Use VertexProvider."""

    def __init__(self, model: str = None):
        logger.warning("OllamaProvider is deprecated. Migrate to VertexProvider.")
        if not HAS_OLLAMA:
            raise ImportError("ollama package not installed")
        self.model = model or os.getenv("OLLAMA_MODEL", "llama2")
        self.client = ollama.Client()

    def generate_json(
        self, prompt: str, schema: Optional[Dict] = None, num_predict: int = 256
    ) -> Dict[str, Any]:
        """Generate JSON using Ollama."""
        try:
            # Add JSON instruction to prompt
            full_prompt = prompt + "\n\nRespond only with valid JSON."

            response = self.client.generate(
                model=self.model,
                prompt=full_prompt,
                format="json" if schema else None,
                options={
                    "num_predict": num_predict,
                    "temperature": 0.3,  # Lower temperature for more consistent output
                },
            )

            # Parse the response
            response_text = response.get("response", "").strip()
            if response_text:
                try:
                    result = json.loads(response_text)
                    return {
                        "value": result,
                        "metrics": {
                            "model": self.model,
                            "provider": "ollama",
                            "tokens": response.get("eval_count", 0),
                        },
                    }
                except json.JSONDecodeError:
                    # Try to extract JSON from response
                    import re

                    json_match = re.search(r"\{[^}]+\}", response_text)
                    if json_match:
                        try:
                            result = json.loads(json_match.group())
                            return {
                                "value": result,
                                "metrics": {"model": self.model, "provider": "ollama"},
                            }
                        except:
                            pass

            return {
                "value": {},
                "metrics": {
                    "model": self.model,
                    "provider": "ollama",
                    "error": "parse_failed",
                },
            }

        except Exception as e:
            return {
                "value": {},
                "metrics": {"model": self.model, "provider": "ollama", "error": str(e)},
            }


class HeuristicProvider(LLMProvider):
    """Fallback heuristic provider when no LLM is available."""

    def generate_json(
        self, prompt: str, schema: Optional[Dict] = None, num_predict: int = 256
    ) -> Dict[str, Any]:
        """Return empty result - heuristics only."""
        return {"value": {}, "metrics": {"provider": "heuristics", "method": "no_llm"}}


_provider_instance = None


def get_llm_provider() -> Optional[LLMProvider]:
    """Get the configured LLM provider instance.

    Defaults to Vertex AI with Gemini 3 Flash. Set LLM_PROVIDER=heuristics to disable.
    """
    global _provider_instance

    if _provider_instance is not None:
        return _provider_instance

    provider_type = os.getenv("LLM_PROVIDER", "vertex").lower()

    # Explicit disable
    if provider_type == "heuristics":
        return None

    # Default: Vertex AI with Gemini 3 Flash
    if provider_type == "vertex" and HAS_VERTEX:
        try:
            _provider_instance = VertexProvider()
            return _provider_instance
        except Exception as e:
            logger.warning("Vertex AI provider init failed: %s", e)

    # Deprecated: Ollama fallback
    if provider_type == "ollama" and HAS_OLLAMA:
        try:
            _provider_instance = OllamaProvider()
            return _provider_instance
        except Exception:
            pass

    # Try Vertex as fallback for unknown providers
    if HAS_VERTEX:
        try:
            _provider_instance = VertexProvider()
            return _provider_instance
        except Exception as e:
            logger.warning("Vertex AI fallback failed: %s", e)

    # Last resort: heuristics only
    _provider_instance = HeuristicProvider()
    return _provider_instance
