"""
Centralized audit logging for all pipelines.
"""

from datetime import datetime, timezone
from typing import Optional

from pyspark.sql import SparkSession
from pyspark.sql.types import (
    StructType,
    StructField,
    TimestampType,
    StringType,
    LongType,
)

from setmore_etl.config import SetmoreConfig


class AuditLogger:
    """Centralized audit logging with explicit schema."""
    
    _AUDIT_SCHEMA = StructType([
        StructField("execution_time", TimestampType(), True),
        StructField("pipeline_name", StringType(), True),
        StructField("extraction_id", StringType(), True),
        StructField("extraction_start", TimestampType(), True),
        StructField("extraction_end", TimestampType(), True),
        StructField("appointments_count", LongType(), True),
        StructField("staff_count", LongType(), True),
        StructField("services_count", LongType(), True),
        StructField("load_mode", StringType(), True),
        StructField("status", StringType(), True),
        StructField("error_message", StringType(), True),
    ])
    
    def __init__(self, spark: SparkSession, config: SetmoreConfig):
        self.spark = spark
        self.config = config
        self._initialized = False
    
    @property
    def table_name(self) -> str:
        """Fully qualified audit table name."""
        return self.config.audit_table_fqn
    
    def _ensure_initialized(self) -> None:
        """Lazy initialization of schema and table."""
        if self._initialized:
            return
        
        self.spark.sql(f"CREATE SCHEMA IF NOT EXISTS {self.config.catalog}.{self.config.schema}")
        
        self.spark.sql(f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                execution_time TIMESTAMP,
                pipeline_name STRING,
                extraction_id STRING,
                extraction_start TIMESTAMP,
                extraction_end TIMESTAMP,
                appointments_count LONG,
                staff_count LONG,
                services_count LONG,
                load_mode STRING,
                status STRING,
                error_message STRING
            ) USING DELTA
        """)
        
        self._initialized = True
    
    def log(
        self,
        pipeline_name: str,
        extraction_id: str,
        window_start: datetime,
        window_end: datetime,
        appointments: int = 0,
        staff: int = 0,
        services: int = 0,
        mode: str = "incremental",
        status: str = "success",
        error: Optional[str] = None,
    ) -> None:
        """Write audit entry with explicit schema to avoid type inference issues."""
        self._ensure_initialized()
        
        data = [(
            datetime.now(timezone.utc),
            pipeline_name,
            extraction_id,
            window_start,
            window_end,
            appointments,
            staff,
            services,
            mode,
            status,
            error,
        )]
        
        df = self.spark.createDataFrame(data, schema=self._AUDIT_SCHEMA)
        df.write.format("delta").mode("append").saveAsTable(self.table_name)