"""Extraction utilities: SQLite state tracking, file I/O helpers.

These are inline equivalents of the extract_core library used in larger
multi-extractor pipelines. For a single-package project they live here;
if you add multiple data sources, extract them to a shared workspace package.
"""

import gzip
import hashlib
import json
import sqlite3
import threading
from pathlib import Path

# ---------------------------------------------------------------------------
# State tracking (SQLite — transactional, stdlib, no extra dependency)
# ---------------------------------------------------------------------------

_CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS extraction_runs (
    run_id        INTEGER PRIMARY KEY AUTOINCREMENT,
    extractor     TEXT NOT NULL,
    started_at    TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')),
    finished_at   TEXT,
    status        TEXT NOT NULL DEFAULT 'running',
    files_written INTEGER DEFAULT 0,
    files_skipped INTEGER DEFAULT 0,
    bytes_written INTEGER DEFAULT 0,
    cursor_value  TEXT,
    error_message TEXT
)
"""


def open_state_db(landing_dir: str | Path) -> sqlite3.Connection:
    """Open (or create) .state.sqlite inside landing_dir.

    WAL mode allows concurrent reads while a run is in progress.
    Caller is responsible for conn.close().
    """
    db_path = Path(landing_dir) / ".state.sqlite"
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(_CREATE_TABLE_SQL)
    conn.commit()
    return conn


def start_run(conn: sqlite3.Connection, extractor: str) -> int:
    """Insert a 'running' row. Returns run_id."""
    cur = conn.execute(
        "INSERT INTO extraction_runs (extractor, status) VALUES (?, 'running')",
        (extractor,),
    )
    conn.commit()
    return cur.lastrowid


def end_run(
    conn: sqlite3.Connection,
    run_id: int,
    *,
    status: str,
    files_written: int = 0,
    files_skipped: int = 0,
    bytes_written: int = 0,
    cursor_value: str | None = None,
    error_message: str | None = None,
) -> None:
    """Update the run row to its final state."""
    assert status in ("success", "failed")
    conn.execute(
        """
        UPDATE extraction_runs
        SET finished_at   = strftime('%Y-%m-%dT%H:%M:%SZ', 'now'),
            status        = ?,
            files_written = ?,
            files_skipped = ?,
            bytes_written = ?,
            cursor_value  = ?,
            error_message = ?
        WHERE run_id = ?
        """,
        (status, files_written, files_skipped, bytes_written, cursor_value, error_message, run_id),
    )
    conn.commit()


def get_last_cursor(conn: sqlite3.Connection, extractor: str) -> str | None:
    """Return the cursor_value from the most recent successful run, or None."""
    row = conn.execute(
        """
        SELECT cursor_value FROM extraction_runs
        WHERE extractor = ? AND status = 'success' AND cursor_value IS NOT NULL
        ORDER BY run_id DESC LIMIT 1
        """,
        (extractor,),
    ).fetchone()
    return row["cursor_value"] if row else None


_SKIP_RESULT = {"files_written": 0, "files_skipped": 1, "bytes_written": 0}


def skip_if_current(conn: sqlite3.Connection, extractor: str, year_month: str) -> dict | None:
    """Return an early-exit result dict if this extractor already ran for year_month.

    Returns None when the extractor should proceed with extraction.
    """
    if get_last_cursor(conn, extractor) == year_month:
        return _SKIP_RESULT
    return None


# ---------------------------------------------------------------------------
# File I/O helpers
# ---------------------------------------------------------------------------


def landing_path(landing_dir: str | Path, *parts: str) -> Path:
    """Return path to a subdirectory of landing_dir, creating it if absent."""
    path = Path(landing_dir).joinpath(*parts)
    path.mkdir(parents=True, exist_ok=True)
    return path


def content_hash(data: bytes, prefix_bytes: int = 8) -> str:
    """SHA256 content fingerprint — used as idempotency key in filenames."""
    assert data, "data must not be empty"
    return hashlib.sha256(data).hexdigest()[:prefix_bytes]


def load_partial_results(partial_path: Path, id_key: str) -> tuple[list[dict], set[str]]:
    """Load already-completed records from a partial JSONL file (crash recovery).

    Returns (records, seen_ids). If the file doesn't exist, returns ([], set()).
    Gracefully handles a truncated last line from a mid-write crash.
    """
    records: list[dict] = []
    seen_ids: set[str] = set()
    if not partial_path.exists():
        return records, seen_ids
    with open(partial_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                records.append(record)
                rid = record.get(id_key)
                if rid:
                    seen_ids.add(rid)
            except json.JSONDecodeError:
                break  # truncated last line from crash — skip it
    return records, seen_ids


def flush_partial_batch(
    partial_file,
    lock: threading.Lock,
    batch: list[dict],
) -> None:
    """Thread-safe batch write of JSON records to the partial JSONL file.

    Writes all records in one lock acquisition with a single flush.
    Call with batches of ~50 records for good I/O throughput vs crash safety tradeoff.
    On crash, at most one batch worth of records is lost.
    """
    assert batch, "batch must not be empty"
    with lock:
        for record in batch:
            partial_file.write(json.dumps(record, separators=(",", ":")) + "\n")
        partial_file.flush()


def write_gzip_atomic(path: Path, data: bytes) -> int:
    """Gzip compress data and write to path atomically via .tmp sibling.

    Returns bytes written. Atomic write means readers never see a partial file.
    """
    assert data, "data must not be empty"
    compressed = gzip.compress(data)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_bytes(compressed)
    tmp.rename(path)
    return len(compressed)


def write_jsonl_atomic(dest: Path, items: list[dict]) -> int:
    """Write items as JSONL, then compress atomically to dest (.jsonl.gz).

    Compresses the working-file → JSONL → gzip pattern into one call.
    Returns compressed bytes written.
    """
    assert items, "items must not be empty"
    working_path = dest.with_suffix(".working.jsonl")
    with open(working_path, "w") as f:
        for item in items:
            f.write(json.dumps(item, separators=(",", ":")) + "\n")
    return compress_jsonl_atomic(working_path, dest)


def compress_jsonl_atomic(jsonl_path: Path, dest_path: Path) -> int:
    """Compress a JSONL working file to .jsonl.gz atomically, then delete the source.

    Streams compression in 1MB chunks (constant memory regardless of file size).
    Atomic via .tmp rename — readers never see a partial .jsonl.gz.
    Deletes the uncompressed working file after successful compression.
    Returns compressed bytes written.
    """
    assert jsonl_path.exists(), f"source must exist: {jsonl_path}"
    assert jsonl_path.stat().st_size > 0, f"source must not be empty: {jsonl_path}"
    tmp = dest_path.with_suffix(dest_path.suffix + ".tmp")
    with open(jsonl_path, "rb") as f_in, gzip.open(tmp, "wb") as f_out:
        while chunk := f_in.read(1_048_576):  # 1 MB chunks
            f_out.write(chunk)
    bytes_written = tmp.stat().st_size
    tmp.rename(dest_path)
    jsonl_path.unlink()
    return bytes_written
