"""
Watermark management for incremental processing with exactly-once semantics.
"""

from datetime import datetime, timezone, timedelta
from typing import Optional, Tuple, Callable

from pyspark.sql import SparkSession

from setmore_etl.config import SetmoreConfig
from setmore_etl.transforms import generate_extraction_id


class WatermarkManager:
    """
    Manages high watermark for incremental processing using actual timestamps.
    
    Tracks precise extraction windows with lookback for late arrivals.
    No side effects on initialization—lazy table creation.
    """
    
    def __init__(
        self,
        spark: SparkSession,
        config: SetmoreConfig,
        clock: Optional[Callable[[], datetime]] = None,
    ):
        self.spark = spark
        self.config = config
        self._clock = clock or (lambda: datetime.now(timezone.utc))
        self._initialized = False
    
    @property
    def table_name(self) -> str:
        """Fully qualified watermark table name."""
        return self.config.watermark_table_fqn
    
    def _ensure_initialized(self) -> None:
        """Lazy initialization of schema and table."""
        if self._initialized:
            return
        
        # Create schema if not exists
        self.spark.sql(f"CREATE SCHEMA IF NOT EXISTS {self.config.catalog}.{self.config.schema}")
        
        # Create watermark table
        self.spark.sql(f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                pipeline_name STRING NOT NULL,
                extraction_id STRING NOT NULL,
                extraction_start TIMESTAMP,
                extraction_end TIMESTAMP,
                records_processed INT,
                run_type STRING,
                status STRING,
                completed_at TIMESTAMP
            ) USING DELTA
        """)
        
        self._initialized = True
    
    def get_last_successful_window(self, pipeline_name: str) -> Optional[Tuple[datetime, datetime]]:
        """Get the actual extraction window of last successful run."""
        self._ensure_initialized()
        
        result = self.spark.sql(f"""
            SELECT extraction_start, extraction_end 
            FROM {self.table_name} 
            WHERE pipeline_name = '{pipeline_name}'
              AND status = 'success'
            ORDER BY completed_at DESC 
            LIMIT 1
        """).collect()
        
        if result and result[0][0] and result[0][1]:
            # Convert naive datetimes from Spark to UTC-aware datetimes
            start = result[0][0]
            end = result[0][1]
            
            # If naive, assume UTC and make aware
            if start.tzinfo is None:
                start = start.replace(tzinfo=timezone.utc)
            if end.tzinfo is None:
                end = end.replace(tzinfo=timezone.utc)
            
            return (start, end)
        return None
    
    def calculate_extraction_window(self, pipeline_name: str) -> Tuple[datetime, datetime, str]:
        """
        Calculate extraction window for incremental processing.
        
        Uses the last successful extraction window's end timestamp as the start
        for the new window. If no watermark exists (first run), uses 
        config.initial_start_date.
        
        Args:
            pipeline_name: Unique identifier for the pipeline
        
        Returns:
            Tuple of (window_start, window_end, extraction_id)
        
        Example:
            First run: window_start = 2000-01-01 (from config)
            Second run: window_start = last extraction_end
        """
        # ✅ Use the existing get_last_successful_window() method
        last_window = self.get_last_successful_window(pipeline_name)
        
        if last_window:
            # Use the END of the last extraction as the START of this one
            extraction_start, extraction_end = last_window
            window_start = extraction_end  # ✅ Start where we left off (now timezone-aware from get_last_successful_window)
            print(f"Found watermark for '{pipeline_name}': last extraction ended at {extraction_end.isoformat()}")
        else:
            # First run: use configured initial date
            window_start = self.config.get_initial_start_datetime()
            # Ensure timezone-aware
            if window_start.tzinfo is None:
                window_start = window_start.replace(tzinfo=timezone.utc)
            print(f"No watermark for '{pipeline_name}'. Using initial_start_date: {window_start.isoformat()}")
        
        window_end = self._clock()  # ✅ Use the injected clock (returns timezone-aware)
        
        # Ensure window_end is also timezone-aware (defensive)
        if window_end.tzinfo is None:
            window_end = window_end.replace(tzinfo=timezone.utc)
        
        extraction_id = generate_extraction_id(window_start, window_end)
        
        return window_start, window_end, extraction_id
    
    def check_extraction_exists(self, extraction_id: str) -> bool:
        """Check idempotency—has this extraction already succeeded?"""
        self._ensure_initialized()
        
        result = self.spark.sql(f"""
            SELECT COUNT(*) as cnt 
            FROM {self.table_name} 
            WHERE extraction_id = '{extraction_id}' 
              AND status = 'success'
        """).collect()
        
        return result[0][0] > 0
    
    def set_watermark(
        self,
        pipeline_name: str,
        extraction_id: str,
        window_start: datetime,
        window_end: datetime,
        records: int,
        run_type: str = "incremental",
        status: str = "success",
    ) -> None:
        """Record extraction window completion."""
        self._ensure_initialized()
        
        self.spark.sql(f"""
            INSERT INTO {self.table_name} 
            VALUES (
                '{pipeline_name}', 
                '{extraction_id}',
                '{window_start}',
                '{window_end}',
                {records},
                '{run_type}',
                '{status}',
                current_timestamp()
            )
        """)