from __future__ import annotations

import datetime
import json
import logging
import os
import time
import re
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
from typing import Any, Dict, List, Optional

from .base import BaseLLMProvider

# Lazy imports for vertexai to avoid slow GCP credential discovery at module load
_vertexai_modules: Dict[str, Any] = {}

def _get_vertexai():
    """Lazy loader for vertexai module."""
    if "vertexai" not in _vertexai_modules:
        try:
            import vertexai
            _vertexai_modules["vertexai"] = vertexai
        except ImportError:
            _vertexai_modules["vertexai"] = None
    return _vertexai_modules["vertexai"]

def _get_generative_model():
    """Lazy loader for GenerativeModel."""
    if "GenerativeModel" not in _vertexai_modules:
        try:
            from vertexai.generative_models import GenerativeModel
            _vertexai_modules["GenerativeModel"] = GenerativeModel
        except ImportError:
            _vertexai_modules["GenerativeModel"] = None
    return _vertexai_modules["GenerativeModel"]

def _get_generation_config():
    """Lazy loader for GenerationConfig."""
    if "GenerationConfig" not in _vertexai_modules:
        try:
            from vertexai.generative_models import GenerationConfig
            _vertexai_modules["GenerationConfig"] = GenerationConfig
        except ImportError:
            _vertexai_modules["GenerationConfig"] = None
    return _vertexai_modules["GenerationConfig"]

def _get_part():
    """Lazy loader for Part."""
    if "Part" not in _vertexai_modules:
        try:
            from vertexai.generative_models import Part
            _vertexai_modules["Part"] = Part
        except ImportError:
            _vertexai_modules["Part"] = None
    return _vertexai_modules["Part"]

def _get_caching():
    """Lazy loader for caching."""
    if "caching" not in _vertexai_modules:
        try:
            from vertexai.preview import caching
            _vertexai_modules["caching"] = caching
        except ImportError:
            _vertexai_modules["caching"] = None
    return _vertexai_modules["caching"]

logger = logging.getLogger(__name__)

# Default fallback model for rate limit / quota errors (GA model with high quotas)
DEFAULT_FALLBACK_MODEL = "gemini-2.5-pro"

# Timeout for Vertex AI generate_content calls (seconds)
# Prevents hanging requests from blocking the API indefinitely
try:
    from fmatch.saas.settings import DEFAULT_VERTEX_TIMEOUT_S
except Exception:
    DEFAULT_VERTEX_TIMEOUT_S = int(
        os.getenv(
            "DEFAULT_VERTEX_TIMEOUT_S",
            os.getenv(
                "DEFAULT_VERTEX_TIMEOUT",
                os.getenv("VERTEX_GENERATE_TIMEOUT", "20"),
            ),
        )
    )


def _vertex_location_for_model(model: Optional[str], location: str) -> str:
    """Gemini 3 preview models require global location."""
    if (model or "").strip().lower().startswith("gemini-3-"):
        return "global"
    return location


def _is_rate_limit_error(exc: Exception) -> bool:
    """Check if exception is a rate limit (429) or quota exceeded error."""
    exc_str = str(exc).lower()
    exc_type = type(exc).__name__.lower()

    # Check for common rate limit indicators
    rate_limit_indicators = [
        "429",
        "rate limit",
        "ratelimit",
        "quota exceeded",
        "resource exhausted",
        "resourceexhausted",
        "too many requests",
    ]

    for indicator in rate_limit_indicators:
        if indicator in exc_str or indicator in exc_type:
            return True

    # Check for google.api_core exceptions
    if hasattr(exc, "code"):
        # gRPC status codes: 8 = RESOURCE_EXHAUSTED, 14 = UNAVAILABLE
        if getattr(exc, "code", None) in (8, 429):
            return True

    return False


def _is_refresh_error(exc: Exception) -> bool:
    """Check if exception is a google.auth RefreshError (credentials expired)."""
    exc_str = str(exc).lower()
    exc_type = type(exc).__name__
    # Check exception type name (handles google.auth.exceptions.RefreshError)
    if "refresherror" in exc_type.lower():
        return True
    # Check exception message for common refresh error indicators
    refresh_indicators = [
        "reauthentication",
        "refresh",
        "token has been expired",
        "invalid_grant",
        "token expired",
        "credentials",
    ]
    for indicator in refresh_indicators:
        if indicator in exc_str:
            return True
    # Check the exception chain
    cause = getattr(exc, "__cause__", None)
    if cause and "refresherror" in type(cause).__name__.lower():
        return True
    return False


class VertexAIProvider(BaseLLMProvider):
    """Vertex AI Gemini provider for structured JSON generations."""

    def __init__(
        self,
        model: Optional[str] = None,
        project: Optional[str] = None,
        location: Optional[str] = None,
    ):
        vertexai = _get_vertexai()
        GenerativeModel = _get_generative_model()
        if vertexai is None or GenerativeModel is None:
            raise ImportError(
                "vertexai is unavailable. Install google-cloud-aiplatform to use the Vertex provider."
            )

        project = (
            project
            or os.getenv("VERTEX_PROJECT")
            or os.getenv("GOOGLE_CLOUD_PROJECT")
            or os.getenv("GCP_PROJECT")
        )
        if not project:
            raise RuntimeError("VERTEX_PROJECT or GOOGLE_CLOUD_PROJECT must be set")

        location = location or os.getenv("VERTEX_LOCATION", "us-central1")
        model = model or os.getenv("VERTEX_MODEL") or os.getenv("GENAI_MODEL") or "gemini-3-flash-preview"
        resolved_location = _vertex_location_for_model(model, location)
        if resolved_location != location:
            logger.info("Vertex AI: overriding location to global for model %s", model)
        location = resolved_location

        super().__init__(model=model)

        # Store for caching support
        self._project = project
        self._location = location
        self._model_name = model

        try:
            vertexai.init(project=project, location=location)
            self._model = GenerativeModel(model)
        except Exception as exc:
            if _is_refresh_error(exc):
                logger.critical(
                    "VERTEX AI AUTH ERROR during initialization: Reauthentication needed. "
                    "Run: gcloud auth application-default login"
                )
                raise RuntimeError(
                    "Reauthentication is needed. Please run "
                    "gcloud auth application-default login to reauthenticate."
                ) from exc
            raise

        # Fallback model for rate limit errors (GA model with high quotas)
        self._fallback_model_name = os.getenv("VERTEX_FALLBACK_MODEL", DEFAULT_FALLBACK_MODEL)
        self._fallback_model: Optional[GenerativeModel] = None
        self._fallback_location = _vertex_location_for_model(self._fallback_model_name, location)

        self._default_config = {
            "temperature": float(os.getenv("VERTEX_TEMPERATURE", "0")),
            "top_p": float(os.getenv("VERTEX_TOP_P", "0.95")),
            "top_k": int(os.getenv("VERTEX_TOP_K", "32")),
            # AI match configs need generous headroom for mappings + reasoning.
            "max_output_tokens": int(os.getenv("VERTEX_MAX_OUTPUT_TOKENS", "8192")),
        }

    def generate_json(
        self,
        prompt: str,
        schema: Optional[Dict[str, Any]] = None,
        cache_name: Optional[str] = None,
        thinking_level: Optional[str] = None,
        **opts: Any,
    ) -> Dict[str, Any]:
        GenerationConfig = _get_generation_config()
        caching = _get_caching()
        GenerativeModel = _get_generative_model()

        config = dict(self._default_config)
        include_thoughts = opts.pop("include_thoughts", None)
        timeout_s = opts.pop("timeout_s", None)
        if include_thoughts is None:
            include_thoughts = os.getenv("VERTEX_INCLUDE_THOUGHTS", "false").lower() in {
                "1",
                "true",
                "yes",
            }
        if timeout_s is None:
            timeout_s = DEFAULT_VERTEX_TIMEOUT_S
        try:
            timeout_s = int(timeout_s)
        except (TypeError, ValueError):
            timeout_s = DEFAULT_VERTEX_TIMEOUT_S
        if "temperature" in opts and opts["temperature"] is not None:
            config["temperature"] = float(opts["temperature"])
        if "max_output_tokens" in opts and opts["max_output_tokens"] is not None:
            config["max_output_tokens"] = int(opts["max_output_tokens"])

        config["response_mime_type"] = "application/json"
        schema_applied = False
        if schema:
            config["response_schema"] = schema
            schema_applied = True

        # Build ThinkingConfig for Gemini 3 models and add to config
        thinking_config_applied = False
        if thinking_level and "gemini-3" in self._model_name.lower():
            try:
                from google.genai import types as genai_types

                valid_levels = {"MINIMAL", "LOW", "MEDIUM", "HIGH"}
                level = thinking_level.upper()
                if level in valid_levels:
                    thinking_config = genai_types.ThinkingConfig(thinking_level=level)
                    config["thinking_config"] = thinking_config
                    logger.debug(
                        "Adding thinking_level=%s for model %s",
                        level,
                        self._model_name,
                    )
                else:
                    logger.warning(
                        "Invalid thinking_level '%s'; valid options are %s",
                        thinking_level,
                        valid_levels,
                    )
            except ImportError:
                logger.debug("ThinkingConfig not available in current SDK version")
            except Exception as exc:
                logger.warning("Failed to create ThinkingConfig: %s", exc)

        if include_thoughts and "gemini-3" in self._model_name.lower():
            config["include_thoughts"] = True

        # Build GenerationConfig, handling schema and thinking_config rejections
        try:
            generation_config = GenerationConfig(**config)
            if "thinking_config" in config:
                thinking_config_applied = True
                logger.debug("Gemini 3 thinking level applied")
        except TypeError as gen_err:
            # SDK may not support thinking_config - try without it
            if "thinking_config" in str(gen_err) and "thinking_config" in config:
                config.pop("thinking_config", None)
                if "include_thoughts" in str(gen_err):
                    config.pop("include_thoughts", None)
                try:
                    generation_config = GenerationConfig(**config)
                    logger.debug("Thinking level not supported by SDK - using default")
                except Exception as inner_exc:
                    if schema_applied:
                        config.pop("response_schema", None)
                        config.pop("include_thoughts", None)
                        generation_config = GenerationConfig(**config)
                    else:
                        raise inner_exc
            elif "include_thoughts" in str(gen_err) and "include_thoughts" in config:
                config.pop("include_thoughts", None)
                try:
                    generation_config = GenerationConfig(**config)
                    logger.debug("include_thoughts not supported by SDK - using default")
                except Exception as inner_exc:
                    if schema_applied:
                        config.pop("response_schema", None)
                        generation_config = GenerationConfig(**config)
                    else:
                        raise inner_exc
            elif schema_applied:
                logger.warning(
                    "Vertex response_schema rejected (%s); retrying without schema", gen_err
                )
                config.pop("response_schema", None)
                config.pop("thinking_config", None)
                config.pop("include_thoughts", None)
                generation_config = GenerationConfig(**config)
            else:
                raise
        except Exception as exc:
            if schema_applied:
                logger.warning(
                    "Vertex response_schema rejected (%s); retrying without schema", exc
                )
                config.pop("response_schema", None)
                config.pop("thinking_config", None)
                config.pop("include_thoughts", None)
                generation_config = GenerationConfig(**config)
            else:
                raise

        # Select model: cached or standard
        model = self._model
        if cache_name and caching:
            try:
                try:
                    cached_content = caching.CachedContent(cache_name=cache_name)
                except TypeError as exc:
                    if "cache_name" in str(exc):
                        cached_content = caching.CachedContent(name=cache_name)
                    else:
                        raise
                model = GenerativeModel.from_cached_content(cached_content=cached_content)
                logger.debug("Using cached context: %s", cache_name)
            except Exception as exc:
                logger.warning(
                    "Cache %s not found or expired (%s). Falling back to standard model.",
                    cache_name,
                    exc,
                )
                model = self._model

        start = time.monotonic()
        used_model = self._model_name
        used_fallback = False

        def _call_generate():
            return model.generate_content([prompt], generation_config=generation_config)

        try:
            # Use timeout to prevent hanging requests from blocking the API
            with ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(_call_generate)
                try:
                    response = future.result(timeout=timeout_s)
                except FuturesTimeoutError:
                    logger.error(
                        "Vertex AI generate_content timed out after %ds for model %s",
                        timeout_s,
                        self._model_name,
                    )
                    raise TimeoutError(
                        f"Vertex AI request timed out after {timeout_s}s"
                    )
        except TimeoutError:
            raise
        except Exception as exc:
            # Check for authentication/refresh errors first - these need admin action
            if _is_refresh_error(exc):
                logger.critical(
                    "VERTEX AI AUTH ERROR: Reauthentication needed. "
                    "Run: gcloud auth application-default login"
                )
                raise RuntimeError(
                    "Reauthentication is needed. Please run "
                    "gcloud auth application-default login to reauthenticate."
                ) from exc
            # Check if this is a rate limit error and we should fallback
            if _is_rate_limit_error(exc) and self._fallback_model_name != self._model_name:
                logger.warning(
                    "Vertex AI rate limit hit on %s, falling back to %s",
                    self._model_name,
                    self._fallback_model_name,
                )
                try:
                    # Lazy-init fallback model (may need different location)
                    if self._fallback_model is None:
                        vertexai = _get_vertexai()
                        if self._fallback_location != self._location:
                            vertexai.init(project=self._project, location=self._fallback_location)
                        self._fallback_model = GenerativeModel(self._fallback_model_name)

                    # Fallback model (e.g., gemini-2.5-pro) may not support thinking_config
                    # so we only pass generation_config for the fallback
                    def _call_fallback():
                        return self._fallback_model.generate_content(
                            [prompt], generation_config=generation_config
                        )

                    with ThreadPoolExecutor(max_workers=1) as executor:
                        future = executor.submit(_call_fallback)
                        try:
                            response = future.result(timeout=timeout_s)
                        except FuturesTimeoutError:
                            raise TimeoutError(
                                f"Vertex AI fallback request timed out after {timeout_s}s"
                            )
                    used_model = self._fallback_model_name
                    used_fallback = True
                except TimeoutError as timeout_exc:
                    logger.error("Vertex AI fallback also timed out: %s", timeout_exc)
                    raise
                except Exception as fallback_exc:
                    logger.exception(
                        "Vertex AI fallback to %s also failed: %s",
                        self._fallback_model_name,
                        fallback_exc,
                    )
                    raise fallback_exc from exc
            else:
                logger.exception("Vertex AI generation failed: %s", exc)
                raise

        elapsed_ms = int((time.monotonic() - start) * 1000)
        if used_fallback:
            logger.info(
                "Vertex AI request completed via fallback model %s in %dms",
                used_model,
                elapsed_ms,
            )

        json_payloads: List[Any] = []
        raw_parts = []
        saw_candidate_parts = False
        thought_snippets = []
        if include_thoughts:
            thought_snippets = self._extract_thoughts(response)
            if thought_snippets:
                preview = " | ".join(thought_snippets)
                logger.debug(
                    "Vertex AI thoughts (%s): %s",
                    used_model,
                    preview[:500],
                )
            else:
                logger.debug("Vertex AI thoughts requested but none returned")
        aggregated_text = getattr(response, "text", None)

        for candidate in getattr(response, "candidates", []) or []:
            finish_reason = getattr(candidate, "finish_reason", None)
            if finish_reason and str(finish_reason).upper() == "MAX_TOKENS":
                logger.warning(
                    "Vertex AI response truncated due to MAX_TOKENS limit (model=%s)",
                    used_model,
                )
            content = getattr(candidate, "content", None)
            if not content:
                continue
            for part in getattr(content, "parts", []) or []:
                json_payload = getattr(part, "json", None)
                if json_payload is not None:
                    json_payloads.append(json_payload)
                    saw_candidate_parts = True
                    continue
                text = getattr(part, "text", None)
                if text:
                    raw_parts.append(text)
                    saw_candidate_parts = True
                    continue
                inline_data = getattr(part, "inline_data", None)
                if inline_data and getattr(inline_data, "data", None):
                    try:
                        raw_parts.append(inline_data.data.decode("utf-8"))
                        saw_candidate_parts = True
                        continue
                    except Exception:
                        pass
                if getattr(part, "mime_type", None) == "application/json":
                    raw_parts.append(getattr(part, "text", "") or "")
                    saw_candidate_parts = True

        aggregated_text_str = ""
        if aggregated_text:
            aggregated_text_str = str(aggregated_text)

        raw_text = ""
        if json_payloads:
            try:
                payload = json_payloads[0] if len(json_payloads) == 1 else json_payloads
                raw_text = json.dumps(payload)
                logger.debug(
                    "Vertex JSON parts detected; using structured payload over raw text."
                )
            except TypeError:
                logger.debug(
                    "Vertex JSON parts not serializable; falling back to raw text."
                )
        if not raw_text:
            raw_from_parts = "".join(raw_parts).strip()
            aggregated_text_clean = aggregated_text_str.strip()
            if aggregated_text_clean:
                if not raw_from_parts or len(raw_from_parts) < len(aggregated_text_clean) * 0.5:
                    raw_text = aggregated_text_clean
                    logger.debug(
                        "Using aggregated Vertex response text (parts length=%d, aggregated length=%d).",
                        len(raw_from_parts),
                        len(aggregated_text_clean),
                    )
                else:
                    raw_text = raw_from_parts
            else:
                raw_text = raw_from_parts
        parsed = self._parse_response_json(raw_text)

        return {
            "value": parsed,
            "provider": "vertex_ai",
            "model": used_model,
            "metrics": {
                "latency_ms": elapsed_ms,
                "fallback_used": used_fallback,
            },
            "raw_text": raw_text,
            "thoughts": thought_snippets,
        }

    @staticmethod
    def _extract_thoughts(response: Any) -> List[str]:
        snippets: List[str] = []
        for candidate in getattr(response, "candidates", []) or []:
            for attr in ("thoughts", "thought"):
                candidate_thought = getattr(candidate, attr, None)
                if candidate_thought:
                    if isinstance(candidate_thought, list):
                        snippets.extend(
                            [str(item).strip() for item in candidate_thought if str(item).strip()]
                        )
                    else:
                        snippets.append(str(candidate_thought).strip())
            content = getattr(candidate, "content", None)
            if not content:
                continue
            for attr in ("thoughts", "thought"):
                content_thought = getattr(content, attr, None)
                if content_thought:
                    if isinstance(content_thought, list):
                        snippets.extend(
                            [str(item).strip() for item in content_thought if str(item).strip()]
                        )
                    else:
                        snippets.append(str(content_thought).strip())
            for part in getattr(content, "parts", []) or []:
                for attr in ("thoughts", "thought"):
                    part_thought = getattr(part, attr, None)
                    if part_thought:
                        if isinstance(part_thought, list):
                            snippets.extend(
                                [str(item).strip() for item in part_thought if str(item).strip()]
                            )
                        else:
                            snippets.append(str(part_thought).strip())
        return [snippet for snippet in snippets if snippet]

    @staticmethod
    def _parse_response_json(raw_text: str) -> Dict[str, Any]:
        """Best-effort JSON parsing that tolerates malformed Gemini output."""
        text = (raw_text or "").strip()
        if not text:
            return {}

        # Strip Markdown-style code fences like ```json ... ```
        if text.startswith("```") and text.endswith("```"):
            text = re.sub(r"^```[a-zA-Z0-9_-]*\s*", "", text)
            text = re.sub(r"\s*```$", "", text)

        attempts = [text]

        # If the model wrapped JSON in extra prose, try extracting the first JSON object.
        start = text.find("{")
        end = text.rfind("}")
        if start != -1 and end != -1 and end > start:
            candidate = text[start : end + 1]
            if candidate != text:
                attempts.append(candidate)

        # Gemini sometimes concatenates multiple JSON objects back-to-back (}{...).
        # Try parsing just the first object in that scenario.
        for splitter in ("}\n{", "}{"):
            idx = text.find(splitter)
            if idx != -1:
                attempts.append(text[: idx + 1])
                break

        # Attempt to repair truncated JSON by balancing brackets
        def repair_truncated(s: str) -> Optional[str]:
            """Balance brackets for truncated JSON responses."""
            if not s or not s.strip():
                return None
            s = s.rstrip()

            # Strategy: Find last position where brackets are balanced after a comma or colon
            # This handles truncation mid-object/array by trimming to last complete element
            best_repair = None

            # Find positions of structural characters from the end
            for i in range(len(s) - 1, -1, -1):
                char = s[i]
                if char in ",:[{":
                    # Try truncating here and balancing
                    candidate = s[:i].rstrip()
                    if candidate.endswith(","):
                        candidate = candidate[:-1]

                    # Check quote balance
                    quote_count = candidate.count('"') - candidate.count('\\"')
                    if quote_count % 2 != 0:
                        continue  # Unclosed string

                    # Balance brackets
                    open_braces = candidate.count("{") - candidate.count("}")
                    open_brackets = candidate.count("[") - candidate.count("]")

                    if open_braces >= 0 and open_brackets >= 0:
                        repaired = candidate + "]" * open_brackets + "}" * open_braces
                        try:
                            json.loads(repaired)
                            best_repair = repaired
                            break
                        except json.JSONDecodeError:
                            continue

            if best_repair:
                return best_repair

            # Fallback: Simple bracket balancing without truncation
            s = re.sub(r",\s*$", "", s)
            open_braces = s.count("{") - s.count("}")
            open_brackets = s.count("[") - s.count("]")
            if open_braces > 0 or open_brackets > 0:
                s += "]" * open_brackets
                s += "}" * open_braces
                try:
                    json.loads(s)
                    return s
                except json.JSONDecodeError:
                    pass
            return None

        for payload in attempts:
            try:
                parsed = json.loads(payload)
            except json.JSONDecodeError:
                # Try repairing truncated JSON
                repaired = repair_truncated(payload)
                if repaired:
                    try:
                        parsed = json.loads(repaired)
                        logger.debug(
                            "Repaired truncated JSON response (added %d brackets)",
                            len(repaired) - len(payload),
                        )
                    except json.JSONDecodeError:
                        continue
                else:
                    continue

            # Sometimes Gemini returns a JSON string; attempt to decode again.
            if isinstance(parsed, str):
                try:
                    parsed = json.loads(parsed)
                except json.JSONDecodeError:
                    logger.warning(
                        "Vertex AI produced a double-encoded JSON string; returning raw text snippet: %s",
                        payload[:120],
                    )
                    return {"raw_text": payload}

            if isinstance(parsed, dict):
                return parsed
            # Wrap non-dict responses under value for consistency.
            return {"value": parsed}

        logger.warning(
            "Vertex AI returned non-JSON content; delivering raw text snippet: %s",
            text[:160],
        )
        return {"raw_text": text}

    def create_cache(
        self,
        content: str,
        ttl_minutes: int = 60,
        display_name: Optional[str] = None,
    ) -> str:
        """
        Upload large context (e.g., Salesforce schema) to Vertex AI and return a cache_name.

        Args:
            content: The content to cache (e.g., schema documentation)
            ttl_minutes: Time-to-live in minutes (default: 60)
            display_name: Optional name for the cache (default: "salesforce_schema_cache")

        Returns:
            The cache name (resource path) or empty string if caching is unavailable
        """
        caching = _get_caching()
        Part = _get_part()
        if caching is None or Part is None:
            logger.warning("Vertex AI Context Caching not available in this SDK version.")
            return ""

        try:
            system_instruction = "You are a Salesforce SOQL expert who generates valid queries from natural language."
            cached = caching.CachedContent.create(
                model_name=self._model_name,
                system_instruction=system_instruction,
                contents=[Part.from_text(content)],
                ttl=datetime.timedelta(minutes=ttl_minutes),
                display_name=display_name or "salesforce_schema_cache",
            )
            logger.info(f"Created Vertex AI Cache: {cached.name} (TTL {ttl_minutes}m)")
            return cached.name
        except Exception as exc:
            logger.error(f"Failed to create cache: {exc}")
            return ""
