"""Auto-generated Time Series Latest Datapoints UDTF for Cognite Data Fusion.

This UDTF retrieves the latest datapoint(s) for one or more Time Series using instance_id.
Uses direct REST calls with scalar execution for Unity Catalog SQL registration compatibility.
Returns rows in format: (time_series_external_id, timestamp, value, status_code)
Note: time_series_external_id is in format "space:external_id" to support time series from different spaces.
"""

from __future__ import annotations

import sys
import traceback
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterator

# NOTE: Helper functions are now defined inside eval() method for Unity Catalog compatibility
# This ensures they're available in Unity Catalog's execution context

from pyspark.sql.types import (
    StructType,
    StructField,
    TimestampType,
    StringType,
    DoubleType,
    IntegerType,
)

{% if use_udtf_decorator %}
from pyspark.sql.functions import udtf

@udtf
{% endif %}
class TimeSeriesLatestDatapointsUDTF:
    """UDTF for retrieving the latest datapoint(s) for one or more Time Series using instance_id.
    
    Uses scalar UDTF execution for Unity Catalog SQL registration compatibility.
    Similar to client.time_series.data.retrieve_latest()
    Returns rows in format: (time_series_external_id, timestamp, value, status_code)
    Note: time_series_external_id is in format "space:external_id" to support time series from different spaces.
    """
    
    def __init__(self) -> None:
        """Initialize UDTF (no parameters allowed when using analyze method).
        
        Token initialization happens in eval() for all registration modes.
        """
        # Initialize instance variables
        self._token_cache: dict[str, str | float] | None = None
        self._init_error: str | None = None
        self._init_error_category: str | None = None
        self._init_success = True
    
    @staticmethod
    def outputSchema() -> StructType:
        """Return schema: (time_series_external_id, timestamp, value, status_code)."""
        fields = [
            StructField("time_series_external_id", StringType(), nullable=True),
            StructField("timestamp", TimestampType(), nullable=True),
            StructField("value", DoubleType(), nullable=True),
            StructField("status_code", IntegerType(), nullable=True),
        ]

        return StructType(fields)
    
    def eval(
        self,
        client_id: str | None = None,
        client_secret: str | None = None,
        tenant_id: str | None = None,
        cdf_cluster: str | None = None,
        project: str | None = None,
        instance_ids: str | None = None,  # Format: "space1:ext_id1,space2:ext_id2"
        before: object | None = None,  # SQL TIMESTAMP, relative time string, ISO 8601
        include_status: object | None = None,  # Boolean
    ) -> Iterator[tuple]:
        """Execute scalar UDTF using instance_ids (comma-separated) for query.
        
        Uses scalar execution for Unity Catalog SQL registration compatibility.
        
        Args:
            client_id: OAuth2 client ID (scalar)
            client_secret: OAuth2 client secret (scalar)
            tenant_id: Azure AD tenant ID (scalar)
            cdf_cluster: CDF cluster URL (scalar)
            project: CDF project name (scalar)
            instance_ids: Comma-separated instance IDs in format "space1:ext_id1,space2:ext_id2" (scalar)
            before: Timestamp before which to retrieve latest datapoints (scalar)
            include_status: Whether to include status information (scalar, boolean)
        
        Yields:
            Tuples with (time_series_external_id, timestamp, value, status_code) columns
        """
        import sys
        
        # CRITICAL: Define helper functions inside eval() for Unity Catalog compatibility
        # Unity Catalog's execution context may not execute module-level code,
        # so all helper functions must be defined here to ensure they're available
        {% set skip_docstring = True %}
        {% filter indent(8, first=True, blank=True) %}
        {% include '_http_client.py.jinja' %}
        {% endfilter %}
        
        try:
            # Helper function to extract scalar from Python input
            def _extract_scalar(value: object | None) -> object | None:
                """Extract scalar value from Python input (scalar mode)."""
                if value is None:
                    return None
                
                # Handle PySpark Column objects (has as_py method)
                if hasattr(value, "as_py"):
                    try:
                        result = value.as_py()
                        if isinstance(result, bytes):
                            return result.decode('utf-8')
                        return result
                    except Exception as e:
                        sys.stderr.write(f"[UDTF] WARNING: as_py() failed for {type(value).__name__}: {e}\n")
                
                # Handle string directly
                if isinstance(value, str):
                    return value
                
                # Handle bytes
                if isinstance(value, bytes):
                    try:
                        return value.decode('utf-8')
                    except UnicodeDecodeError:
                        sys.stderr.write(f"[UDTF] WARNING: Failed to decode bytes value\n")
                        return None
                
                # Try to convert to string as fallback
                try:
                    str_value = str(value)
                    if str_value and str_value.lower() != "none":
                        return str_value
                    return None
                except Exception:
                    return None
            
            # Helper function to extract boolean from Python input
            def _extract_bool(value: object | None) -> bool:
                """Extract boolean value from Python input (scalar mode)."""
                if value is None:
                    return False
                
                # Handle PySpark Column objects (has as_py method)
                if hasattr(value, "as_py"):
                    try:
                        result = value.as_py()
                        if isinstance(result, bool):
                            return result
                        # Convert string representations
                        if isinstance(result, str):
                            return result.lower() in ("true", "1", "yes", "on")
                        # Convert numeric
                        if isinstance(result, (int, float)):
                            return bool(result)
                    except Exception:
                        pass
                
                # Handle boolean directly
                if isinstance(value, bool):
                    return value
                
                # Handle string representations
                if isinstance(value, str):
                    return value.lower() in ("true", "1", "yes", "on")
                
                # Handle numeric
                if isinstance(value, (int, float)):
                    return bool(value)
                
                # Default to False
                return False
            
            # Extract scalar values from Python inputs
            client_id_str = _extract_scalar(client_id)
            client_secret_str = _extract_scalar(client_secret)
            tenant_id_str = _extract_scalar(tenant_id)
            cdf_cluster_str = _extract_scalar(cdf_cluster)
            project_str = _extract_scalar(project)
            
            # Validate required parameters
            if not all([client_id_str, client_secret_str, tenant_id_str, cdf_cluster_str, project_str]):
                error_msg = "Missing required parameters: client_id, client_secret, tenant_id, cdf_cluster, or project"
                error_category = 'CONFIGURATION'
                sys.stderr.write(f"[UDTF] ✗ {error_category}: {error_msg}\n")

                yield (None, None, None, None)

                return
            
            instance_ids_str = _extract_scalar(instance_ids)
            if not instance_ids_str:
                error_msg = "instance_ids is required (format: 'space1:ext_id1,space2:ext_id2')"
                error_category = 'CONFIGURATION'
                sys.stderr.write(f"[UDTF] ✗ {error_category}: {error_msg}\n")

                yield (None, None, None, None)

                return
            
            # Convert before parameter to string (supports SQL TIMESTAMP, relative time, ISO 8601)
            def timestamp_to_string(timestamp: object) -> str:
                """Convert timestamp to string format accepted by CDF API.
                
                Supports:
                - SQL TIMESTAMP types (datetime objects from Unity Catalog) -> ISO 8601
                - Relative time strings: "1h-ago", "2w-ago", "now"
                - ISO 8601 strings: "2024-01-01T00:00:00Z"
                - SQL timestamp strings: "2024-01-01 00:00:00" -> ISO 8601
                """
                import re
                from datetime import datetime, timezone
                
                # Handle None - default to "now"
                if timestamp is None:
                    return "now"
                
                # Handle datetime objects (from SQL TIMESTAMP) -> convert to ISO 8601
                if hasattr(timestamp, 'timestamp') or hasattr(timestamp, 'isoformat'):
                    try:
                        # datetime.datetime or similar object
                        dt = timestamp
                        # Ensure it's timezone-aware and in UTC
                        if hasattr(dt, 'tzinfo'):
                            if dt.tzinfo is None:
                                dt = dt.replace(tzinfo=timezone.utc)
                            else:
                                dt = dt.astimezone(timezone.utc)
                        # Convert to ISO 8601 format
                        if hasattr(dt, 'isoformat'):
                            return dt.isoformat(timespec='milliseconds').replace('+00:00', 'Z')
                        else:
                            return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"  # Include milliseconds, add Z
                    except (AttributeError, TypeError):
                        pass
                
                # Also check for datetime-like objects with year/month/day attributes
                if hasattr(timestamp, 'year') and hasattr(timestamp, 'month') and hasattr(timestamp, 'day'):
                    try:
                        # Try to construct a datetime and convert
                        if hasattr(timestamp, 'hour'):
                            dt = datetime(
                                timestamp.year, timestamp.month, timestamp.day,
                                timestamp.hour, timestamp.minute, timestamp.second,
                                getattr(timestamp, 'microsecond', 0)
                            )
                        else:
                            dt = datetime(timestamp.year, timestamp.month, timestamp.day)
                        # Assume UTC if no timezone info
                        if not hasattr(timestamp, 'tzinfo') or timestamp.tzinfo is None:
                            dt = dt.replace(tzinfo=timezone.utc)
                        return dt.isoformat(timespec='milliseconds').replace('+00:00', 'Z')
                    except Exception:
                        pass
                
                # Handle string types - pass through (CDF API accepts relative time and ISO 8601)
                if isinstance(timestamp, str):
                    # If it's already a relative time string or ISO 8601, return as-is
                    if timestamp.lower() == "now":
                        return "now"
                    
                    # Check if it's a relative time string
                    match = re.match(r"(\d+)([smhdw])-(ago|ahead)", timestamp.lower())
                    if match:
                        return timestamp  # Return relative time string as-is
                    
                    # Check if it's SQL timestamp format (YYYY-MM-DD HH:MM:SS) -> convert to ISO 8601
                    try:
                        for fmt in [
                            "%Y-%m-%d %H:%M:%S.%f",
                            "%Y-%m-%d %H:%M:%S",
                        ]:
                            try:
                                dt = datetime.strptime(timestamp, fmt)
                                # Convert to ISO 8601 format
                                return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
                            except ValueError:
                                continue
                    except Exception:
                        pass
                    
                    # If it looks like ISO 8601, return as-is
                    if "T" in timestamp or timestamp.endswith("Z"):
                        return timestamp
                    
                    # Otherwise, try to parse and convert
                    try:
                        for fmt in [
                            "%Y-%m-%dT%H:%M:%S.%fZ",
                            "%Y-%m-%dT%H:%M:%SZ",
                            "%Y-%m-%dT%H:%M:%S",
                        ]:
                            try:
                                dt = datetime.strptime(timestamp, fmt)
                                return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
                            except ValueError:
                                continue
                    except Exception:
                        pass
                    
                    # If we can't parse it, return as-is (might be valid for CDF API)
                    return timestamp
                
                # For other types, try to convert to string
                return str(timestamp)
            
            before_str = timestamp_to_string(before) or "now"
            include_status_bool = _extract_bool(include_status)
            
            # Parse instance_ids
            node_ids = []
            for instance_id_str in instance_ids_str.split(","):
                instance_id_str = instance_id_str.strip()
                if not instance_id_str:
                    continue
                if ":" not in instance_id_str:
                    error_msg = f"Invalid instance_id format '{instance_id_str}'. Expected format: 'space:external_id'"
                    error_category = 'CONFIGURATION'
                    sys.stderr.write(f"[UDTF] ✗ {error_category}: {error_msg}\n")

                    yield (None, None, None, None)

                    return
                space, external_id = instance_id_str.split(":", 1)
                space = space.strip()
                external_id = external_id.strip()
                if not space or not external_id:
                    error_msg = f"Invalid instance_id format '{instance_id_str}'. Both space and external_id must be non-empty."
                    error_category = 'CONFIGURATION'
                    sys.stderr.write(f"[UDTF] ✗ {error_category}: {error_msg}\n")

                    yield (None, None, None, None)

                    return
                node_ids.append({"space": space, "externalId": external_id})
            
            if not node_ids:
                error_msg = "At least one valid instance_id is required"
                error_category = 'CONFIGURATION'
                sys.stderr.write(f"[UDTF] ✗ {error_category}: {error_msg}\n")

                yield (None, None, None, None)

                return
            
            # Get OAuth token (with caching)
            try:
                token, expires_at = _get_oauth_token(
                    str(client_id_str),
                    str(client_secret_str),
                    str(tenant_id_str),
                    str(cdf_cluster_str),
                    self._token_cache,
                )
                # Cache token for next call
                self._token_cache = {
                    "access_token": token,
                    "expires_at": expires_at,
                }
            except Exception as token_error:
                error_category = _classify_error(None, token_error)
                sys.stderr.write(f"[UDTF] ✗ {error_category}: Failed to acquire OAuth token: {token_error}\n")

                yield (None, None, None, None)

                return
            
            try:
                # Log query start (both stderr and stdout for Serverless SQL visibility)
                query_start = f"[UDTF] Querying latest datapoints for {len(node_ids)} time series"
                sys.stderr.write(f"{query_start}\n")
                print(query_start)  # Also print to stdout for Serverless SQL visibility
                
                # Make direct REST call to CDF API
                base_url = f"https://{cdf_cluster_str}.cognitedata.com"
                api_path = f"/api/v1/projects/{project_str}/timeseries/data/latest"
                full_url = f"{base_url}{api_path}"
                
                # Build request payload
                payload = {
                    "items": [
                        {
                            "instanceId": {
                                "space": node_id["space"],
                                "externalId": node_id["externalId"],
                            }
                        }
                        for node_id in node_ids
                    ],
                    "before": before_str,
                    "includeStatus": include_status_bool,
                }
                
                # Make request
                response = _make_request(
                    method="POST",
                    url=full_url,
                    token=token,
                    json_data=payload,
                    timeout=30,
                )
                
                response_data = response.json()
                items = response_data.get("items", [])
                
                # Import datetime and timezone for timestamp conversion
                from datetime import datetime, timezone
                
                row_count = 0
                for item in items:
                    instance_id = item.get("instanceId", {})
                    space_val = instance_id.get("space", "")
                    external_id_val = instance_id.get("externalId", "")
                    ts_external_id = f"{space_val}:{external_id_val}"
                    
                    # Get the latest datapoint
                    datapoint = item.get("datapoint")
                    if datapoint:
                        ts_ms = datapoint.get("timestamp", 0)
                        timestamp_dt = datetime.fromtimestamp(ts_ms / 1000, tz=timezone.utc) if ts_ms else None
                        val = datapoint.get("value")
                        status = datapoint.get("status", {})
                        status_code = status.get("code") if include_status_bool else None
                        
                        yield (ts_external_id, timestamp_dt, val, status_code)
                        row_count += 1
                
                # Log success (both stderr and stdout for Serverless SQL visibility)
                success_msg = f"[UDTF] SUCCESS: yielded {row_count} rows (scalar mode)"
                sys.stderr.write(f"{success_msg}\n")
                print(success_msg)  # Also print to stdout for Serverless SQL visibility

            except Exception as e:
                # Log error for debugging (without traceback to avoid import issues)
                # Both stderr and stdout for Serverless SQL visibility
                error_info = f"[UDTF] ✗ Error during query: {type(e).__name__}: {str(e)}"
                sys.stderr.write(error_info)
                print(error_info)  # Also print to stdout for Serverless SQL visibility

                yield (None, None, None, None)

                return
                
        except Exception as outer_error:
            # Last resort: if anything goes wrong, yield empty row
            # Both stderr and stdout for Serverless SQL visibility
            error_info = f"ERROR: Unexpected error in eval(): {type(outer_error).__name__}: {str(outer_error)}"
            sys.stderr.write(f"{error_info}\n")
            print(error_info)  # Also print to stdout for Serverless SQL visibility

            yield (None, None, None, None)

            return

