use std::{future::Future, time::Instant};

use arrow_array::RecordBatch;
use tracing::Instrument as _;

use crate::{
    Result, WritePhase,
    observability::{
        BATCH_WRITE_COMPLETED_EVENT, BATCH_WRITE_FAILED_EVENT, BATCH_WRITE_PHASE, BATCH_WRITE_SPAN,
        BATCH_WRITE_STARTED_EVENT, TRACE_TARGET, duration_micros_u64, usize_to_u64_saturating,
    },
    write::writer::{WriteBackend, WriteStats},
};

use super::{backend_trace_name, diagnostic_codes_for_error, sanitized_error_summary};

#[derive(Debug)]
pub(crate) struct BatchWriteTrace {
    span: tracing::Span,
    started: Instant,
    backend: WriteBackend,
    attempted_batch_ordinal: u64,
    batch_row_count: u64,
    batch_column_count: u64,
    batch_is_empty: bool,
    accepted_rows_before: u64,
    accepted_batches_before: u64,
}

impl BatchWriteTrace {
    pub(crate) fn new(backend: WriteBackend, stats: WriteStats, batch: &RecordBatch) -> Self {
        let attempted_batch_ordinal = stats.batches_written.saturating_add(1);
        let batch_row_count = usize_to_u64_saturating(batch.num_rows());
        let batch_column_count = usize_to_u64_saturating(batch.num_columns());
        let batch_is_empty = batch.num_rows() == 0;
        let span = tracing::info_span!(
            target: TRACE_TARGET,
            BATCH_WRITE_SPAN,
            phase = BATCH_WRITE_PHASE,
            backend = backend_trace_name(backend),
            attempted_batch_ordinal,
            batch_row_count,
            batch_column_count,
            batch_is_empty,
            accepted_rows_before = stats.rows_written,
            accepted_batches_before = stats.batches_written,
        );

        Self {
            span,
            started: Instant::now(),
            backend,
            attempted_batch_ordinal,
            batch_row_count,
            batch_column_count,
            batch_is_empty,
            accepted_rows_before: stats.rows_written,
            accepted_batches_before: stats.batches_written,
        }
    }

    pub(crate) fn emit_started(&self) {
        let _span_guard = self.span.enter();
        tracing::info!(
            target: TRACE_TARGET,
            phase = BATCH_WRITE_PHASE,
            telemetry_event = BATCH_WRITE_STARTED_EVENT,
            backend = backend_trace_name(self.backend),
            attempted_batch_ordinal = self.attempted_batch_ordinal,
            batch_row_count = self.batch_row_count,
            batch_column_count = self.batch_column_count,
            batch_is_empty = self.batch_is_empty,
            accepted_rows_before = self.accepted_rows_before,
            accepted_batches_before = self.accepted_batches_before,
            batch_write_result = "started"
        );
    }

    /// Emits batch lifecycle events while polling the write future inside this span.
    pub(crate) async fn trace_result<F>(&self, future: F) -> Result<WriteStats>
    where
        F: Future<Output = Result<WriteStats>>,
    {
        self.emit_started();
        let result = future.instrument(self.span.clone()).await;
        match &result {
            Ok(stats) => self.emit_completed(*stats),
            Err(err) => {
                let phase = err.write_phase().unwrap_or(WritePhase::BatchWrite);
                self.emit_failed(phase.as_str(), err);
            }
        }
        result
    }

    pub(crate) fn emit_completed(&self, stats: WriteStats) {
        let _span_guard = self.span.enter();
        tracing::info!(
            target: TRACE_TARGET,
            phase = BATCH_WRITE_PHASE,
            telemetry_event = BATCH_WRITE_COMPLETED_EVENT,
            backend = backend_trace_name(self.backend),
            attempted_batch_ordinal = self.attempted_batch_ordinal,
            batch_row_count = self.batch_row_count,
            batch_column_count = self.batch_column_count,
            batch_is_empty = self.batch_is_empty,
            accepted_rows_before = self.accepted_rows_before,
            accepted_batches_before = self.accepted_batches_before,
            accepted_rows_after = stats.rows_written,
            accepted_batches_after = stats.batches_written,
            batch_write_result = "success",
            elapsed_us = duration_micros_u64(self.started.elapsed())
        );
    }

    pub(crate) fn emit_failed(&self, phase: &'static str, error: &crate::Error) {
        let diagnostic_codes = diagnostic_codes_for_error(error);
        let _span_guard = self.span.enter();
        tracing::error!(
            target: TRACE_TARGET,
            phase,
            telemetry_event = BATCH_WRITE_FAILED_EVENT,
            backend = backend_trace_name(self.backend),
            attempted_batch_ordinal = self.attempted_batch_ordinal,
            batch_row_count = self.batch_row_count,
            batch_column_count = self.batch_column_count,
            batch_is_empty = self.batch_is_empty,
            accepted_rows_before = self.accepted_rows_before,
            accepted_batches_before = self.accepted_batches_before,
            batch_write_result = "failure",
            error_summary = sanitized_error_summary(error),
            diagnostic_codes = diagnostic_codes.as_str(),
            elapsed_us = duration_micros_u64(self.started.elapsed())
        );
    }
}

#[cfg(test)]
mod tests {
    use std::{
        borrow::Cow,
        future::Future,
        sync::Arc,
        task::{Context, Poll, Waker},
    };

    use arrow_array::{Int32Array, RecordBatch, StringArray};
    use arrow_schema::{DataType, Field, Schema};

    use super::BatchWriteTrace;
    use crate::{
        Diagnostic, DiagnosticCode, DiagnosticSet, Error,
        observability::{
            BATCH_SCHEMA_VALIDATION_PHASE, BATCH_WRITE_COMPLETED_EVENT, BATCH_WRITE_FAILED_EVENT,
            BATCH_WRITE_PHASE, BATCH_WRITE_SPAN, PACKET_WRITE_PHASE, VALUE_CONVERSION_PHASE,
            test_support::{assert_trace_field, capture_traces, trace_event},
        },
        write::writer::{WriteBackend, WriteStats},
    };

    #[test]
    fn baseline_batch_write_success_emits_stats_trace() -> Result<(), String> {
        let batch = int32_batch("id", &[10, 20]);
        let stats_before = WriteStats::default();
        let stats_after = WriteStats {
            rows_written: 2,
            batches_written: 1,
        };

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(WriteBackend::BaselineTokenRow, stats_before, &batch);
            trace.emit_started();
            trace.emit_completed(stats_after);
        });

        let records = traces.records()?;
        let event = trace_event(&records, BATCH_WRITE_COMPLETED_EVENT)?;
        assert_eq!(event.span_name(), Some(BATCH_WRITE_SPAN));
        assert_trace_field(event, "phase", BATCH_WRITE_PHASE);
        assert_trace_field(event, "backend", "BaselineTokenRow");
        assert_trace_field(event, "attempted_batch_ordinal", "1");
        assert_trace_field(event, "batch_row_count", "2");
        assert_trace_field(event, "batch_column_count", "1");
        assert_trace_field(event, "batch_is_empty", "false");
        assert_trace_field(event, "accepted_rows_before", "0");
        assert_trace_field(event, "accepted_batches_before", "0");
        assert_trace_field(event, "accepted_rows_after", "2");
        assert_trace_field(event, "accepted_batches_after", "1");
        assert_trace_field(event, "batch_write_result", "success");
        assert!(event.fields().contains_key("elapsed_us"));

        Ok(())
    }

    #[test]
    fn empty_batch_write_success_trace_matches_stats() -> Result<(), String> {
        let batch = int32_batch("id", &[]);
        let stats_after = WriteStats {
            rows_written: 0,
            batches_written: 1,
        };

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(
                WriteBackend::BaselineTokenRow,
                WriteStats::default(),
                &batch,
            );
            trace.emit_completed(stats_after);
        });

        let records = traces.records()?;
        let event = trace_event(&records, BATCH_WRITE_COMPLETED_EVENT)?;
        assert_trace_field(event, "batch_row_count", "0");
        assert_trace_field(event, "batch_is_empty", "true");
        assert_trace_field(event, "accepted_rows_after", "0");
        assert_trace_field(event, "accepted_batches_after", "1");

        Ok(())
    }

    #[test]
    fn batch_schema_validation_failure_emits_trace_without_accepting_batch() -> Result<(), String> {
        let batch = int32_batch("id", &[1]);
        let error = value_conversion_error(DiagnosticCode::SchemaMismatch);

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(
                WriteBackend::BaselineTokenRow,
                WriteStats::default(),
                &batch,
            );
            trace.emit_failed(BATCH_SCHEMA_VALIDATION_PHASE, &error);
        });

        let records = traces.records()?;
        let event = trace_event(&records, BATCH_WRITE_FAILED_EVENT)?;
        assert_eq!(event.span_name(), Some(BATCH_WRITE_SPAN));
        assert_trace_field(event, "phase", BATCH_SCHEMA_VALIDATION_PHASE);
        assert_trace_field(event, "backend", "BaselineTokenRow");
        assert_trace_field(event, "batch_row_count", "1");
        assert_trace_field(event, "batch_column_count", "1");
        assert_trace_field(event, "accepted_rows_before", "0");
        assert_trace_field(event, "accepted_batches_before", "0");
        assert_trace_field(event, "batch_write_result", "failure");
        assert_trace_field(event, "diagnostic_codes", "SchemaMismatch");

        Ok(())
    }

    #[test]
    fn value_conversion_failure_emits_trace_without_accepting_batch() -> Result<(), String> {
        let batch = int32_batch("amount", &[1, 2]);
        let error = value_conversion_error(DiagnosticCode::NonFiniteFloat);

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(
                WriteBackend::BaselineTokenRow,
                WriteStats::default(),
                &batch,
            );
            trace.emit_failed(VALUE_CONVERSION_PHASE, &error);
        });

        let records = traces.records()?;
        let event = trace_event(&records, BATCH_WRITE_FAILED_EVENT)?;
        assert_trace_field(event, "phase", VALUE_CONVERSION_PHASE);
        assert_trace_field(event, "diagnostic_codes", "NonFiniteFloat");
        assert_trace_field(
            event,
            "error_summary",
            "value conversion failed with diagnostics",
        );

        Ok(())
    }

    #[test]
    fn packet_write_failure_emits_trace_without_accepting_batch() -> Result<(), String> {
        let batch = int32_batch("id", &[1, 2, 3]);
        let error = Error::Tiberius {
            source: tiberius::error::Error::BulkInput(Cow::Borrowed("fake send failure")),
        };

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(
                WriteBackend::BaselineTokenRow,
                WriteStats::default(),
                &batch,
            );
            trace.emit_failed(PACKET_WRITE_PHASE, &error);
        });

        let records = traces.records()?;
        let event = trace_event(&records, BATCH_WRITE_FAILED_EVENT)?;
        assert_trace_field(event, "phase", PACKET_WRITE_PHASE);
        assert_trace_field(event, "backend", "BaselineTokenRow");
        assert_trace_field(event, "diagnostic_codes", "");
        assert_trace_field(event, "error_summary", "tiberius operation failed");

        Ok(())
    }

    #[test]
    fn batch_write_trace_does_not_emit_row_values() -> Result<(), String> {
        let batch = utf8_batch("secret_value", &["password=secret"]);

        let (_result, traces) = capture_traces(|| {
            let trace = BatchWriteTrace::new(
                WriteBackend::BaselineTokenRow,
                WriteStats::default(),
                &batch,
            );
            trace.emit_completed(WriteStats {
                rows_written: 1,
                batches_written: 1,
            });
        });

        traces.assert_no_forbidden_text(&["password=secret"])?;

        Ok(())
    }

    #[test]
    fn tiberius_protocol_events_are_parented_to_batch_write_span() -> Result<(), String> {
        const TIBERIUS_PROTOCOL_TARGET: &str = "tiberius_raw_bulk::protocol";
        const TIBERIUS_PACKET_EVENT: &str = "protocol.bulk_load.packet.written";

        let batch = int32_batch("id", &[10, 20]);
        let (result, traces) = capture_traces(|| {
            let trace =
                BatchWriteTrace::new(WriteBackend::DirectRawBulk, WriteStats::default(), &batch);
            let write_result = poll_ready(trace.trace_result(async {
                tracing::info!(
                    target: TIBERIUS_PROTOCOL_TARGET,
                    telemetry_event = TIBERIUS_PACKET_EVENT,
                    phase = PACKET_WRITE_PHASE,
                    operation = "bulk_load",
                    status = "success",
                    packet_count = 1_u64
                );
                Ok(WriteStats {
                    rows_written: 2,
                    batches_written: 1,
                })
            }))?;
            write_result.map_err(|err| err.to_string())?;
            Ok::<(), String>(())
        });
        result?;

        let records = traces.records()?;
        let event = trace_event(&records, TIBERIUS_PACKET_EVENT)?;
        assert_eq!(event.target(), TIBERIUS_PROTOCOL_TARGET);
        assert_eq!(event.span_name(), Some(BATCH_WRITE_SPAN));
        assert_trace_field(event, "phase", PACKET_WRITE_PHASE);
        assert_trace_field(event, "operation", "bulk_load");
        assert_trace_field(event, "status", "success");
        assert_trace_field(event, "packet_count", "1");

        Ok(())
    }

    #[test]
    fn direct_batch_write_success_emits_stats_trace() -> Result<(), String> {
        let batch = int32_batch("id", &[10, 20]);

        for backend in [WriteBackend::DirectFramedBulk, WriteBackend::DirectRawBulk] {
            let (_result, traces) = capture_traces(|| {
                let trace = BatchWriteTrace::new(backend, WriteStats::default(), &batch);
                trace.emit_completed(WriteStats {
                    rows_written: 2,
                    batches_written: 1,
                });
            });

            let records = traces.records()?;
            let event = trace_event(&records, BATCH_WRITE_COMPLETED_EVENT)?;
            let expected_backend = match backend {
                WriteBackend::DirectFramedBulk => "DirectFramedBulk",
                WriteBackend::DirectRawBulk => "DirectRawBulk",
                WriteBackend::Auto => "Auto",
                WriteBackend::BaselineTokenRow => "BaselineTokenRow",
            };
            assert_trace_field(event, "backend", expected_backend);
            assert_trace_field(event, "batch_row_count", "2");
            assert_trace_field(event, "batch_column_count", "1");
            assert_trace_field(event, "accepted_rows_after", "2");
            assert_trace_field(event, "accepted_batches_after", "1");
        }

        Ok(())
    }

    #[test]
    fn direct_raw_value_conversion_failure_trace_includes_diagnostic_codes() -> Result<(), String> {
        let batch = int32_batch("u64_value", &[1]);
        let error = value_conversion_error(DiagnosticCode::IntegerOutOfRange);

        let (_result, traces) = capture_traces(|| {
            let trace =
                BatchWriteTrace::new(WriteBackend::DirectRawBulk, WriteStats::default(), &batch);
            trace.emit_failed(VALUE_CONVERSION_PHASE, &error);
        });

        let records = traces.records()?;
        let failure = trace_event(&records, BATCH_WRITE_FAILED_EVENT)?;
        assert_eq!(failure.span_name(), Some(BATCH_WRITE_SPAN));
        assert_trace_field(failure, "phase", VALUE_CONVERSION_PHASE);
        assert_trace_field(failure, "backend", "DirectRawBulk");
        assert_trace_field(failure, "batch_row_count", "1");
        assert_trace_field(failure, "batch_column_count", "1");
        assert_trace_field(failure, "diagnostic_codes", "IntegerOutOfRange");
        assert_trace_field(
            failure,
            "error_summary",
            "value conversion failed with diagnostics",
        );

        Ok(())
    }

    fn value_conversion_error(code: DiagnosticCode) -> Error {
        Error::ValueConversion {
            diagnostics: DiagnosticSet::from(vec![Diagnostic::error(code, "test diagnostic")]),
        }
    }

    fn int32_batch(name: &str, values: &[i32]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)]));
        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values.to_vec()))]).unwrap()
    }

    fn utf8_batch(name: &str, values: &[&str]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Utf8, false)]));
        RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(values.to_vec()))]).unwrap()
    }

    fn poll_ready<F>(future: F) -> Result<F::Output, String>
    where
        F: Future,
    {
        let mut future = Box::pin(future);
        let mut context = Context::from_waker(Waker::noop());
        match future.as_mut().poll(&mut context) {
            Poll::Ready(output) => Ok(output),
            Poll::Pending => Err("test future unexpectedly pending".to_owned()),
        }
    }
}
