from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ingestion.common.sql import sql_string

# Incremental strategies for external (non-extracted) object-store sources. A strategy
# decides three things for a bucket of objects: how to split them into loadable
# *increments*, what cursor value identifies each (the delta key), and which provenance
# columns record it. All strategies normalise the cursor to epoch milliseconds and expose
# a single BIGINT watermark column, so the orchestration's delta logic
# (`key > max(watermark)`) is identical regardless of strategy.


@dataclass(frozen=True)
class Increment:
    """One unit loaded together: a cursor ``key`` (epoch ms, the delta key) and the object
    names that make it up. ``label`` is for run events/logs."""

    key: int
    label: str
    files: list[str]


class IncrementalStrategy:
    name: str

    def provenance_columns(self) -> list[tuple[str, str]]:
        """Strategy-specific raw columns (``(name, databend_type)``)."""
        raise NotImplementedError

    @property
    def watermark_column(self) -> str:
        """The BIGINT epoch-ms column the loader maxes for the delta check."""
        raise NotImplementedError

    def increments(self, objects: list[dict[str, Any]]) -> list[Increment]:
        """Split discovered objects (``[{name, size, last_modified_ms}]``) into increments."""
        raise NotImplementedError

    def clear_predicate(self, increment: Increment) -> str:
        """SQL boolean selecting this increment's already-present rows (for idempotent reload)."""
        raise NotImplementedError

    def stamp_assignments(self, increment: Increment) -> dict[str, str]:
        """Column -> SQL literal stamped onto this increment's freshly-copied rows. Must use
        the same value(s) as ``clear_predicate`` so reload/count target the same rows."""
        raise NotImplementedError


class FilenameStrategy(IncrementalStrategy):
    """Cursor is an epoch-ms timestamp captured from each object name; objects sharing it
    are one snapshot, loaded together. Used by metadb (``run-<ts>-part-…parquet``)."""

    name = "filename"

    def __init__(self, pattern: str) -> None:
        self._pattern = re.compile(pattern)

    def provenance_columns(self) -> list[tuple[str, str]]:
        return [("_extraction_ts", "BIGINT"), ("_extraction_at", "TIMESTAMP")]

    @property
    def watermark_column(self) -> str:
        return "_extraction_ts"

    def increments(self, objects: list[dict[str, Any]]) -> list[Increment]:
        grouped: dict[int, list[str]] = {}
        for item in objects:
            name = item["name"]
            if not name.endswith(".parquet"):
                continue
            match = self._pattern.search(name)
            if match:
                grouped.setdefault(int(match.group("ts")), []).append(name)
        return [Increment(key=ts, label=str(ts), files=sorted(files)) for ts, files in grouped.items()]

    def clear_predicate(self, increment: Increment) -> str:
        return f"_extraction_ts = {increment.key}"

    def stamp_assignments(self, increment: Increment) -> dict[str, str]:
        return {"_extraction_ts": str(increment.key), "_extraction_at": f"to_timestamp({increment.key}, 3)"}


class LastModifiedStrategy(IncrementalStrategy):
    """Cursor is each object's last-modified time (a dltHub filesystem-incremental shape):
    each file is its own increment, loaded when its mtime is newer than the max loaded."""

    name = "last_modified"

    def provenance_columns(self) -> list[tuple[str, str]]:
        return [
            ("_source_file", "STRING"),
            ("_source_last_modified_ms", "BIGINT"),
            ("_source_last_modified", "TIMESTAMP"),
        ]

    @property
    def watermark_column(self) -> str:
        return "_source_last_modified_ms"

    def increments(self, objects: list[dict[str, Any]]) -> list[Increment]:
        increments = []
        for item in objects:
            name = item["name"]
            if not name.endswith(".parquet"):
                continue
            ms = item.get("last_modified_ms")
            if ms is None:
                raise RuntimeError(f"object {name} has no last_modified; cannot use last_modified strategy")
            increments.append(Increment(key=int(ms), label=Path(name).name, files=[name]))
        return increments

    def clear_predicate(self, increment: Increment) -> str:
        return f"_source_file = {sql_string(increment.files[0])}"

    def stamp_assignments(self, increment: Increment) -> dict[str, str]:
        # _source_file stores the full object key (matches clear_predicate); the COPY uses the
        # stage-relative path separately.
        return {
            "_source_file": sql_string(increment.files[0]),
            "_source_last_modified_ms": str(increment.key),
            "_source_last_modified": f"to_timestamp({increment.key}, 3)",
        }


def build_strategy(incremental: dict[str, Any]) -> IncrementalStrategy:
    """Construct the strategy named in a source's ``incremental`` config block."""
    name = incremental.get("strategy", "filename")
    if name == "filename":
        return FilenameStrategy(incremental.get("pattern", r"run-(?P<ts>\d+)-"))
    if name == "last_modified":
        return LastModifiedStrategy()
    raise ValueError(f"unknown incremental strategy: {name!r} (expected 'filename' or 'last_modified')")


__all__ = [
    "Increment",
    "IncrementalStrategy",
    "FilenameStrategy",
    "LastModifiedStrategy",
    "build_strategy",
]