"""
Unit tests for DuckDB Executors.

Tests sync and async executors with mocked filesystems.
"""

import pytest
from unittest.mock import MagicMock, patch, AsyncMock
import asyncio
import json

from owilix.core.db import (
    OWIDuckDBSelectExecutor,
    OWIDuckDBAsyncExecutor,
    OWIlixSQLQuery,
    ParquetBatch,
    ParquetBatchResult
)


class TestOWIlixSQLQuery:
    """Tests for SQL query builder."""
    
    def test_from_templates(self):
        sql = OWIlixSQLQuery.from_templates("pq_select")
        assert "SELECT" in sql.sql
        
    def test_files(self):
        sql = OWIlixSQLQuery("SELECT * FROM read_parquet(${owi_remote_files})")
        result = sql.files(["/path/file1.parquet", "/path/file2.parquet"])
        assert "/path/file1.parquet" in result.sql
        assert "/path/file2.parquet" in result.sql
        
    def test_where(self):
        sql = OWIlixSQLQuery("SELECT * FROM t ${where}")
        result = sql.where("col = 'value'")
        assert "WHERE col = 'value'" in result.sql
        
    def test_limit(self):
        sql = OWIlixSQLQuery("SELECT * FROM t ${limit}")
        result = sql.limit(100)
        assert "LIMIT 100" in result.sql
        
    def test_chaining(self):
        sql = (OWIlixSQLQuery.from_templates("pq_select")
               .select("url,title")
               .where("url LIKE '%example%'"))
        assert "url,title" in sql.sql
        assert "example" in sql.sql


class TestParquetBatch:
    """Tests for ParquetBatch model."""
    
    def test_creation(self):
        batch = ParquetBatch(
            files=[("/path/file.parquet", "")],
            query_args={"key": "value"}
        )
        assert len(batch.files) == 1
        assert batch.query_args["key"] == "value"


class TestParquetBatchResult:
    """Tests for ParquetBatchResult model."""
    
    def test_success_result(self):
        batch = ParquetBatch(files=[("/path/file.parquet", "")], query_args={})
        result = ParquetBatchResult(parquet_batch=batch, rows=[{"a": 1}], success=True)
        assert result.success is True
        assert len(result.rows) == 1
        
    def test_error_result(self):
        batch = ParquetBatch(files=[("/path/file.parquet", "")], query_args={})
        error = Exception("Test error")
        result = ParquetBatchResult(parquet_batch=batch, success=False, error=error)
        assert result.success is False
        assert result.error is not None


class TestOWIDuckDBSelectExecutor:
    """Tests for sync executor."""
    
    @pytest.fixture
    def mock_fs(self):
        fs = MagicMock()
        fs.protocol = "http2irods"
        fs.fsid = "http2irods"
        return fs
    
    @pytest.fixture
    def sample_sql(self):
        return OWIlixSQLQuery("SELECT COUNT(*) FROM read_parquet(${owi_remote_files})")
    
    def test_init(self, mock_fs, sample_sql):
        pq_files = {mock_fs: [("/path/file.parquet", "")]}
        executor = OWIDuckDBSelectExecutor(
            pq_files=pq_files,
            owilix_sql=sample_sql,
            pq_batch_size=10,
            prefetch=2
        )
        assert executor.pq_batch_size == 10
        assert executor.prefetch == 2
        asyncio.run(executor.close())
        
    def test_generate_ordered_tasks(self, mock_fs, sample_sql):
        files = [(f"/path/file{i}.parquet", "") for i in range(25)]
        pq_files = {mock_fs: files}
        executor = OWIDuckDBSelectExecutor(
            pq_files=pq_files,
            owilix_sql=sample_sql,
            pq_batch_size=10
        )
        tasks = executor.generate_ordered_tasks()
        # 25 files / 10 per batch = 3 batches
        assert len(tasks) == 3
        asyncio.run(executor.close())
        
    def test_default_batch_size_is_10(self, mock_fs, sample_sql):
        pq_files = {mock_fs: [("/path/file.parquet", "")]}
        executor = OWIDuckDBSelectExecutor(
            pq_files=pq_files,
            owilix_sql=sample_sql
        )
        assert executor.pq_batch_size == 10
        asyncio.run(executor.close())

    @patch("owilix.core.db.duckdb_executor.DuckDBConnectionPool")
    def test_run_query_logs_sql_when_enabled(self, mock_pool_cls, mock_fs, sample_sql, tmp_path):
        from owilix.core.sql_logging import configure_sql_logging

        log_path = tmp_path / "sql.log"
        configure_sql_logging(True, str(log_path))

        mock_cursor = MagicMock()
        mock_cursor.fetchmany.side_effect = [[(1,)], []]
        mock_cursor.description = [("count_star()",)]
        mock_conn = MagicMock()
        mock_conn.filesystem_is_registered.return_value = True
        mock_conn.execute.return_value = mock_cursor
        mock_pool = MagicMock()
        mock_pool.acquire_connection.return_value = (mock_conn, MagicMock())
        mock_pool_cls.return_value = mock_pool

        pq_files = {mock_fs: [("/path/file.parquet", "")]}
        executor = OWIDuckDBSelectExecutor(pq_files=pq_files, owilix_sql=sample_sql)
        batch = executor.generate_ordered_tasks()[0][1]
        results = list(executor.run_query_on_batch(mock_fs, batch))

        assert results[0].success is True
        entries = [json.loads(line) for line in log_path.read_text(encoding="utf-8").splitlines() if line.strip()]
        assert any("read_parquet" in entry["sql"] for entry in entries)

        configure_sql_logging(False, None)


class TestOWIDuckDBAsyncExecutor:
    """Tests for async executor."""
    
    @pytest.fixture
    def mock_fs(self):
        fs = MagicMock()
        fs.protocol = "http2irods"
        fs.fsid = "http2irods"
        return fs
    
    @pytest.fixture
    def sample_sql(self):
        return OWIlixSQLQuery("SELECT COUNT(*) FROM read_parquet(${owi_remote_files})")
    
    def test_init(self, mock_fs, sample_sql):
        pq_files = {mock_fs: [("/path/file.parquet", "")]}
        executor = OWIDuckDBAsyncExecutor(
            pq_files=pq_files,
            owilix_sql=sample_sql,
            url_base="https://example.com/irods",
            token="test-token",
            pq_batch_size=10,
            max_concurrent=5
        )
        assert executor.pq_batch_size == 10
        assert executor.max_concurrent == 5
        # Clean up (sync because we haven't entered async context)
        asyncio.run(executor.close())
        
    def test_generate_ordered_tasks(self, mock_fs, sample_sql):
        files = [(f"/path/file{i}.parquet", "") for i in range(25)]
        pq_files = {mock_fs: files}
        executor = OWIDuckDBAsyncExecutor(
            pq_files=pq_files,
            owilix_sql=sample_sql,
            url_base="https://example.com/irods",
            token="test-token",
            pq_batch_size=10
        )
        tasks = executor.generate_ordered_tasks()
        # 25 files / 10 per batch = 3 batches
        assert len(tasks) == 3
        asyncio.run(executor.close())


class TestDuckDBConnectionPool:
    """Tests for DuckDBConnectionPool acquire/release/close lifecycle."""

    @patch("owilix.core.db.duckdb_executor.duckdb")
    @patch("owilix.core.db.duckdb_executor.tempfile")
    def test_pool_init_size(self, mock_tempfile, mock_duckdb):
        mock_tmp = MagicMock()
        mock_tmp.name = "/tmp/test"
        mock_tempfile.TemporaryDirectory.return_value = mock_tmp
        mock_duckdb.connect.return_value = MagicMock()

        from owilix.core.db.duckdb_executor import DuckDBConnectionPool
        pool = DuckDBConnectionPool(size=3, memory_limit="4GB")
        assert len(pool._pool) == 3

    @patch("owilix.core.db.duckdb_executor.duckdb")
    @patch("owilix.core.db.duckdb_executor.tempfile")
    def test_acquire_release(self, mock_tempfile, mock_duckdb):
        mock_tmp = MagicMock()
        mock_tmp.name = "/tmp/test"
        mock_tempfile.TemporaryDirectory.return_value = mock_tmp
        mock_duckdb.connect.return_value = MagicMock()

        from owilix.core.db.duckdb_executor import DuckDBConnectionPool
        pool = DuckDBConnectionPool(size=2)
        assert len(pool._pool) == 2

        conn_tuple = pool.acquire_connection()
        assert len(pool._pool) == 1

        pool.release_connection(conn_tuple)
        assert len(pool._pool) == 2

    @patch("owilix.core.db.duckdb_executor.duckdb")
    @patch("owilix.core.db.duckdb_executor.tempfile")
    def test_close_all(self, mock_tempfile, mock_duckdb):
        mock_conn = MagicMock()
        mock_tmp = MagicMock()
        mock_tmp.name = "/tmp/test"
        mock_tempfile.TemporaryDirectory.return_value = mock_tmp
        mock_duckdb.connect.return_value = mock_conn

        from owilix.core.db.duckdb_executor import DuckDBConnectionPool
        pool = DuckDBConnectionPool(size=2)
        pool.close_all()

        assert len(pool._pool) == 0
        assert mock_conn.close.call_count == 2
        assert mock_tmp.cleanup.call_count == 2


class TestRetryQuery:
    """Tests for _retry_query logic in OWIDuckDBSelectExecutor."""

    @pytest.fixture
    def executor_with_mock_pool(self):
        mock_fs = MagicMock()
        mock_fs.protocol = "file"
        mock_fs.fsid = "file"
        with patch("owilix.core.db.duckdb_executor.DuckDBConnectionPool"):
            executor = OWIDuckDBSelectExecutor(
                pq_files={mock_fs: [("/f.pq", "")]},
                owilix_sql=OWIlixSQLQuery("SELECT 1"),
                retry_count=3,
            )
        return executor

    def test_retry_on_connection_error(self, executor_with_mock_pool):
        executor = executor_with_mock_pool
        fn = MagicMock(side_effect=[ConnectionError("fail"), "ok"])
        with patch.object(executor, "_handle_retry_delay"):
            result = executor._retry_query(fn)
        assert result == "ok"
        assert fn.call_count == 2

    def test_retry_exhausted_raises(self, executor_with_mock_pool):
        executor = executor_with_mock_pool
        fn = MagicMock(side_effect=ConnectionError("fail"))
        with patch.object(executor, "_handle_retry_delay"):
            with pytest.raises(ConnectionError):
                executor._retry_query(fn)
        assert fn.call_count == 3

    def test_no_retry_on_unexpected_error(self, executor_with_mock_pool):
        executor = executor_with_mock_pool
        fn = MagicMock(side_effect=ValueError("bad"))
        with pytest.raises(ValueError):
            executor._retry_query(fn)
        assert fn.call_count == 1

    def test_stop_iteration_returns_none(self, executor_with_mock_pool):
        executor = executor_with_mock_pool
        fn = MagicMock(side_effect=StopIteration())
        result = executor._retry_query(fn)
        assert result is None


@pytest.mark.integration
class TestExecutorIntegration:
    """Integration tests requiring network (marked for skip in CI)."""
    
    @pytest.mark.skip(reason="Requires LEXIS connection")
    def test_sync_executor_real_query(self):
        """Test sync executor with real data."""
        pass
        
    @pytest.mark.skip(reason="Requires LEXIS connection")  
    def test_async_executor_real_query(self):
        """Test async executor with real data."""
        pass