from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field
from pathlib import Path
import multiprocessing as mp
import queue
import time
from typing import Any, Iterable

from .metrics import maybe_report, start_throughput
from .sink import (
    CSVSink,
    CollectorSink,
    CountSink,
    NDJSONSink,
    ParquetSink,
    RecordSink,
    StreamCSVSink,
    StreamJSONSink,
)
from . import uopy_adapter
from .import_utils import call_factory, load_callable

logger = logging.getLogger("universe_reads")


_SENTINEL = None  # picklable across processes


@dataclass(frozen=True)
class RunOptions:
    sessions: int
    count_only: bool
    output_dir: Path | None
    batch_size: int
    progress_every_seconds: float
    max_records: int | None
    dry_run: int | None
    session_factory_kwargs: dict[str, Any]
    dry_run_data: dict[str, Any] = field(default_factory=dict)
    retries: int = 3
    retry_backoff: float = 0.5
    output_format: str = "ndjson"        # ndjson | jsonl | csv | parquet
    stream: bool = False                 # stream to stdout instead of files
    where_clause: str | None = None      # e.g. "WITH STATUS = 'ACTIVE'"
    select_override: str | None = None   # full SELECT statement override
    partition: str | None = None         # "prefix" | "hash" | None


@dataclass(frozen=True)
class WorkerResult:
    worker_id: int
    records: int
    errors: int
    elapsed_seconds: float
    data: dict[str, Any] = field(default_factory=dict)


def _make_sink(*, worker_id: int, options: RunOptions, stream_queue: Any | None = None) -> RecordSink:
    """Create the appropriate sink based on RunOptions."""
    if options.count_only:
        return CountSink()

    # Streaming mode: push to a mp.Queue that the main process drains to stdout.
    if options.stream and stream_queue is not None:
        return _QueueSink(stream_queue)

    if options.output_dir is not None:
        fmt = options.output_format.lower()
        ext = {"ndjson": "ndjson", "jsonl": "jsonl", "csv": "csv", "parquet": "parquet"}.get(fmt, "ndjson")
        path = options.output_dir / f"worker-{worker_id}.{ext}"
        if fmt in ("ndjson", "jsonl"):
            return NDJSONSink(path)
        elif fmt == "csv":
            return CSVSink(path)
        elif fmt == "parquet":
            return ParquetSink(path)

    return CollectorSink()


class _QueueSink:
    """Push records into a multiprocessing queue for the main process to drain."""

    def __init__(self, q: Any) -> None:
        self._q = q
        self.count = 0
        self.data: dict[str, Any] = {}  # not used, but keeps interface consistent

    def on_record(self, *, key: str, record: Any) -> None:
        self._q.put((key, record))
        self.count += 1

    def close(self) -> None:
        pass


def _iter_keys_dry_run(n: int, data: dict[str, Any] | None = None) -> Iterable[str]:
    if data:
        yield from list(data.keys())[:n]
        return
    for i in range(n):
        yield f"KEY-{i:09d}"


def _read_record_dry_run(key: str, data: dict[str, Any] | None = None) -> Any:
    if data and key in data:
        return data[key]
    return f"RECORD({key})"


def _close_session(session: Any) -> None:
    """Safely close a uopy session, ignoring errors."""
    try:
        close = getattr(session, "close", None)
        if callable(close):
            close()
    except Exception:
        pass


def _create_session(
    session_factory_spec: str,
    kwargs: dict[str, Any],
) -> Any:
    """Import the factory and create a fresh session."""
    factory = load_callable(session_factory_spec)
    return call_factory(factory, kwargs)


def _reconnect_session(
    *,
    worker_id: int,
    old_session: Any | None,
    session_factory_spec: str,
    kwargs: dict[str, Any],
) -> Any:
    """Close the old session and create a new one."""
    logger.warning("worker-%d: reconnecting session after RPC failure", worker_id)
    if old_session is not None:
        _close_session(old_session)
        # Clear the cached file handle so the next read opens a fresh one.
        if hasattr(old_session, "_universe_reads_file_cache"):
            delattr(old_session, "_universe_reads_file_cache")
    return _create_session(session_factory_spec, kwargs)


def _is_connection_error(exc: BaseException) -> bool:
    """Heuristic: is this likely a transport / RPC-level failure?

    If the uopy package is available we check for its UOError base class.
    Otherwise we fall back to message sniffing.
    """
    try:
        import uopy  # type: ignore
        if isinstance(exc, uopy.UOError):
            return True
    except Exception:
        pass
    msg = str(exc).lower()
    return any(tok in msg for tok in ("rpc", "connection", "socket", "timeout", "broken pipe", "eof"))


def _worker_loop(
    *,
    worker_id: int,
    file_name: str,
    session_factory_spec: str | None,
    key_queue: Any,
    result_queue: Any,
    options: RunOptions,
    stream_queue: Any | None = None,
) -> None:
    started = time.monotonic()
    errors = 0
    retries_total = 0
    throughput = start_throughput()

    sink = _make_sink(worker_id=worker_id, options=options, stream_queue=stream_queue)

    session: Any | None = None
    if options.dry_run is None:
        if not session_factory_spec:
            raise RuntimeError("session_factory_spec is required when not in dry-run")
        session = _create_session(session_factory_spec, options.session_factory_kwargs)

    max_retries = max(0, options.retries)
    backoff_base = max(0.0, options.retry_backoff)

    def _read_batch(keys: list[str]) -> None:
        """Try to read a batch of keys, with retries + session reconnection."""
        nonlocal session, errors, retries_total

        if options.dry_run is not None:
            for key in keys:
                sink.on_record(key=key, record=_read_record_dry_run(key, options.dry_run_data or None))
                throughput.total += 1
            return

        last_exc: BaseException | None = None
        for attempt in range(1 + max_retries):
            try:
                records, batch_errors = uopy_adapter.read_records(session, file_name, keys)
                errors += batch_errors
                for key, record in zip(keys, records, strict=False):
                    sink.on_record(key=key, record=record)
                    throughput.total += 1
                return  # success
            except Exception as exc:
                last_exc = exc
                retries_total += 1
                if attempt < max_retries:
                    delay = backoff_base * (2 ** attempt)
                    logger.warning(
                        "worker-%d: batch of %d failed (attempt %d/%d): %s — retrying in %.1fs",
                        worker_id, len(keys), attempt + 1, 1 + max_retries, exc, delay,
                    )
                    time.sleep(delay)
                    # Reconnect if it looks like a transport error.
                    if _is_connection_error(exc) and session_factory_spec:
                        session = _reconnect_session(
                            worker_id=worker_id,
                            old_session=session,
                            session_factory_spec=session_factory_spec,
                            kwargs=options.session_factory_kwargs,
                        )

        # All retries exhausted.
        errors += len(keys)
        logger.error(
            "worker-%d: batch of %d FAILED after %d attempts: %s — dead-letter keys: %s",
            worker_id, len(keys), 1 + max_retries, last_exc,
            ", ".join(keys[:20]) + (" ..." if len(keys) > 20 else ""),
        )

    def _read_single(key: str) -> None:
        """Try to read a single key, with retries + session reconnection."""
        nonlocal session, errors, retries_total

        if options.dry_run is not None:
            record = _read_record_dry_run(key, options.dry_run_data or None)
            sink.on_record(key=key, record=record)
            throughput.total += 1
            return

        last_exc: BaseException | None = None
        for attempt in range(1 + max_retries):
            try:
                record = uopy_adapter.read_record(session, file_name, key)
                sink.on_record(key=key, record=record)
                throughput.total += 1
                return  # success
            except Exception as exc:
                last_exc = exc
                retries_total += 1
                if attempt < max_retries:
                    delay = backoff_base * (2 ** attempt)
                    logger.warning(
                        "worker-%d: key %s failed (attempt %d/%d): %s — retrying in %.1fs",
                        worker_id, key, attempt + 1, 1 + max_retries, exc, delay,
                    )
                    time.sleep(delay)
                    if _is_connection_error(exc) and session_factory_spec:
                        session = _reconnect_session(
                            worker_id=worker_id,
                            old_session=session,
                            session_factory_spec=session_factory_spec,
                            kwargs=options.session_factory_kwargs,
                        )

        errors += 1
        logger.error(
            "worker-%d: key %s FAILED after %d attempts: %s",
            worker_id, key, 1 + max_retries, last_exc,
        )

    try:
        while True:
            try:
                batch = key_queue.get(timeout=1.0)
            except queue.Empty:
                continue

            if batch is _SENTINEL:
                break

            if isinstance(batch, list):
                _read_batch(batch)
            else:
                _read_single(batch)

            msg = maybe_report(
                throughput,
                label=f"worker-{worker_id}",
                every_seconds=options.progress_every_seconds,
            )
            if msg:
                logger.info(msg)

            if options.max_records is not None and throughput.total >= options.max_records:
                break

    finally:
        sink.close()
        _close_session(session)
        elapsed = time.monotonic() - started
        if retries_total > 0:
            logger.info("worker-%d: %d retry attempts during run", worker_id, retries_total)
        result_queue.put(
            WorkerResult(
                worker_id=worker_id,
                records=getattr(sink, "count", 0),
                errors=errors,
                elapsed_seconds=elapsed,
                data=getattr(sink, "data", {}),
            )
        )


def _shard_key(key: str, sessions: int, partition: str | None) -> int:
    """Determine which worker shard a key belongs to."""
    if partition == "prefix":
        ordinal = ord(key[0]) if key else 0
        return ordinal % sessions
    if partition == "hash":
        return hash(key) % sessions
    return 0


def run_parallel_read(
    *,
    file_name: str,
    session_factory_spec: str | None,
    options: RunOptions,
) -> list[WorkerResult]:
    """Read records using N concurrent sessions (default 5).

    Strategy:
    - Main process is both producer and worker-0 (it reads its own shard).
    - (N-1) additional worker processes each open their own UniVerse session.
    - Keys are sharded by modulo (default), prefix, or hash.
    - In streaming mode records flow through a mp.Queue to stdout.
    """

    if options.sessions < 1:
        raise ValueError("sessions must be >= 1")

    sessions = int(options.sessions)
    use_partition = options.partition in ("prefix", "hash")

    ctx = mp.get_context("spawn")

    # Streaming: shared queue all workers push (key, record) into.
    stream_queue: Any | None = None
    if options.stream:
        stream_queue = ctx.Queue(maxsize=options.batch_size * sessions * 2)

    # Dedicated queue per worker avoids contention.
    key_queues = [ctx.Queue(maxsize=options.batch_size * 4) for _ in range(sessions)]
    result_queue = ctx.Queue()

    procs: list[mp.Process] = []
    for worker_id in range(1, sessions):
        p = ctx.Process(
            target=_worker_loop,
            kwargs={
                "worker_id": worker_id,
                "file_name": file_name,
                "session_factory_spec": session_factory_spec,
                "key_queue": key_queues[worker_id],
                "result_queue": result_queue,
                "options": options,
                "stream_queue": stream_queue,
            },
            name=f"uv-worker-{worker_id}",
        )
        p.daemon = True
        p.start()
        procs.append(p)

    # Start the streaming drain thread if applicable.
    _drain_thread: threading.Thread | None = None
    if options.stream and stream_queue is not None:
        fmt = options.output_format.lower()
        if fmt == "csv":
            _stream_sink: Any = StreamCSVSink()
        else:
            _stream_sink = StreamJSONSink()

        def _drain() -> None:
            while True:
                try:
                    item = stream_queue.get(timeout=2.0)
                except Exception:
                    continue
                if item is _SENTINEL:
                    break
                key, record = item
                _stream_sink.on_record(key=key, record=record)
            _stream_sink.close()

        _drain_thread = threading.Thread(target=_drain, name="stream-drain", daemon=True)
        _drain_thread.start()

    # Main process: open its own session and do key enumeration.
    started = time.monotonic()
    throughput = start_throughput()
    errors = 0

    sink0 = _make_sink(worker_id=0, options=options, stream_queue=stream_queue)
    session0: Any | None = None
    if options.dry_run is None:
        if not session_factory_spec:
            raise RuntimeError("session_factory_spec is required when not in dry-run")
        session_factory = load_callable(session_factory_spec)
        session0 = call_factory(session_factory, options.session_factory_kwargs)

    def put_batch(worker_id: int, batch: list[str]) -> None:
        # Blocks if worker is slow; that’s fine and provides backpressure.
        key_queues[worker_id].put(batch)

    batch_buffers: list[list[str]] = [[] for _ in range(sessions)]

    if options.dry_run is not None:
        dry_data = options.dry_run_data or None
        keys_iter: Iterable[str] = _iter_keys_dry_run(options.dry_run, dry_data)
    else:
        # Filtered key iteration when WHERE or SELECT override is specified.
        if options.where_clause or options.select_override:
            keys_iter = uopy_adapter.iter_keys_filtered(
                session0,
                file_name,
                where_clause=options.where_clause,
                select_override=options.select_override,
            )
        else:
            keys_iter = uopy_adapter.iter_keys(session0, file_name)

    try:
        buf0: list[str] = []

        def flush_buf0() -> None:
            nonlocal errors
            if not buf0:
                return
            try:
                if options.dry_run is not None:
                    for k in buf0:
                        sink0.on_record(key=k, record=_read_record_dry_run(k, dry_data))
                        throughput.total += 1
                else:
                    records, batch_errors = uopy_adapter.read_records(session0, file_name, buf0)
                    errors += batch_errors
                    for k, r in zip(buf0, records, strict=False):
                        sink0.on_record(key=k, record=r)
                        throughput.total += 1
            except Exception:
                errors += len(buf0)
            buf0.clear()

        for idx, key in enumerate(keys_iter):
            if use_partition:
                shard = _shard_key(key, sessions, options.partition)
            else:
                shard = idx % sessions

            if shard == 0:
                buf0.append(key)
                if len(buf0) >= options.batch_size:
                    flush_buf0()

                msg = maybe_report(
                    throughput,
                    label="main",
                    every_seconds=options.progress_every_seconds,
                )
                if msg:
                    logger.info(msg)

            else:
                buf = batch_buffers[shard]
                buf.append(key)
                if len(buf) >= options.batch_size:
                    put_batch(shard, buf)
                    batch_buffers[shard] = []

            if options.max_records is not None:
                if throughput.total >= options.max_records:
                    break

        # Flush remaining key batches
        flush_buf0()
        for worker_id in range(1, sessions):
            buf = batch_buffers[worker_id]
            if buf:
                put_batch(worker_id, buf)

    finally:
        # Signal workers to stop
        for worker_id in range(1, sessions):
            key_queues[worker_id].put(_SENTINEL)

        sink0.close()
        try:
            close = getattr(session0, "close", None)
            if callable(close):
                close()
        except Exception:
            pass

    # Collect results
    results: list[WorkerResult] = []

    # Add main as a WorkerResult too
    results.append(
        WorkerResult(
            worker_id=0,
            records=getattr(sink0, "count", 0),
            errors=errors,
            elapsed_seconds=time.monotonic() - started,
            data=getattr(sink0, "data", {}),
        )
    )

    # Wait for workers + gather their results (with guardrails against deadlocks)
    deadline = time.monotonic() + 60.0
    remaining = sessions - 1
    while remaining > 0 and time.monotonic() < deadline:
        try:
            res = result_queue.get(timeout=1.0)
        except queue.Empty:
            continue
        if isinstance(res, WorkerResult):
            results.append(res)
            remaining -= 1

    # If any workers didn't report, terminate them so we don't hang.
    if remaining > 0:
        for p in procs:
            if p.is_alive():
                p.terminate()

    for p in procs:
        p.join(timeout=5)

    # Signal the stream drain thread to finish and wait for it.
    if stream_queue is not None:
        stream_queue.put(_SENTINEL)
    if _drain_thread is not None:
        _drain_thread.join(timeout=10)

    return sorted(results, key=lambda r: r.worker_id)
