import pytest
import pandas as pd
from unittest.mock import Mock, patch, MagicMock
from CrumblPy.timescale import TimescaleToolKit


class TestTimescaleToolKit:

    def test_init_with_explicit_credentials(self):
        ts = TimescaleToolKit(
            user='test_user',
            password='test_password',
            host='localhost',
            port=5433,
            database='testdb'
        )
        assert ts.user == 'test_user'
        assert ts.password == 'test_password'
        assert ts.host == 'localhost'
        assert ts.port == 5433
        assert ts.database == 'testdb'
        assert ts.engine is None

    @patch.dict('os.environ', {
        'TIMESCALE_USER': 'env_user',
        'TIMESCALE_PASSWORD': 'env_password',
    }, clear=False)
    def test_init_with_env_vars(self):
        ts = TimescaleToolKit()
        assert ts.user == 'env_user'
        assert ts.password == 'env_password'
        assert ts.host == 'stats_timescale_primary.crumbl.dev'
        assert ts.port == 5432
        assert ts.database == 'crumblanalytics'

    @patch.dict('os.environ', {
        'TIMESCALE_USER': 'env_user',
        'TIMESCALE_PASSWORD': 'env_password',
        'TIMESCALE_HOST': 'env_host',
        'TIMESCALE_PORT': '9999',
        'TIMESCALE_DATABASE': 'env_db',
    }, clear=False)
    def test_init_with_all_env_vars(self):
        ts = TimescaleToolKit()
        assert ts.host == 'env_host'
        assert ts.port == 9999
        assert ts.database == 'env_db'

    @patch('CrumblPy.timescale.create_engine')
    def test_connect(self, mock_create_engine):
        mock_engine = Mock()
        mock_create_engine.return_value = mock_engine

        ts = TimescaleToolKit(
            user='u',
            password='p',
            host='h',
            port=5432,
            database='d'
        )
        result = ts.connect()

        assert result == mock_engine
        assert ts.engine == mock_engine
        mock_create_engine.assert_called_once_with(
            'postgresql://u:p@h:5432/d'
        )

    @patch('pandas.read_sql')
    @patch('CrumblPy.timescale.create_engine')
    def test_fetch_data(self, mock_create_engine, mock_read_sql):
        mock_engine = Mock()
        mock_create_engine.return_value = mock_engine
        mock_df = pd.DataFrame({'col1': [1, 2], 'col2': ['a', 'b']})
        mock_read_sql.return_value = mock_df

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        result = ts.fetch_data("SELECT * FROM test_table")

        assert isinstance(result, pd.DataFrame)
        assert result.equals(mock_df)
        mock_read_sql.assert_called_once_with("SELECT * FROM test_table", mock_engine)
        assert ts.engine == mock_engine

    @patch('pandas.read_sql')
    @patch('CrumblPy.timescale.create_engine')
    def test_fetch_data_auto_connect(self, mock_create_engine, mock_read_sql):
        mock_engine = Mock()
        mock_create_engine.return_value = mock_engine
        mock_df = pd.DataFrame({'col1': [1, 2]})
        mock_read_sql.return_value = mock_df

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        assert ts.engine is None
        result = ts.fetch_data("SELECT 1")

        mock_create_engine.assert_called_once()
        mock_read_sql.assert_called_once_with("SELECT 1", mock_engine)

    @patch('pandas.DataFrame.to_sql')
    @patch('CrumblPy.timescale.create_engine')
    def test_insert_data(self, mock_create_engine, mock_to_sql):
        mock_engine = Mock()
        mock_create_engine.return_value = mock_engine

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        test_df = pd.DataFrame({'col1': [1, 2], 'col2': ['a', 'b']})
        ts.insert_data(test_df, 'test_table')

        mock_to_sql.assert_called_once_with(
            'test_table', mock_engine, if_exists='append', index=False
        )

    @patch('CrumblPy.timescale.create_engine')
    def test_execute_query(self, mock_create_engine):
        mock_engine = Mock()
        mock_conn = MagicMock()
        mock_engine.connect.return_value.__enter__ = Mock(return_value=mock_conn)
        mock_engine.connect.return_value.__exit__ = Mock(return_value=None)
        mock_create_engine.return_value = mock_engine

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        ts.execute_query("UPDATE test_table SET col1 = 1")

        mock_conn.execute.assert_called_once()
        mock_conn.commit.assert_called_once()

    @patch('CrumblPy.timescale.create_engine')
    def test_execute_query_auto_connect(self, mock_create_engine):
        mock_engine = Mock()
        mock_conn = MagicMock()
        mock_engine.connect.return_value.__enter__ = Mock(return_value=mock_conn)
        mock_engine.connect.return_value.__exit__ = Mock(return_value=None)
        mock_create_engine.return_value = mock_engine

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        ts.engine = None
        ts.execute_query("DELETE FROM test_table")

        mock_create_engine.assert_called_once()
        mock_conn.execute.assert_called_once()
        mock_conn.commit.assert_called_once()

    @patch('CrumblPy.timescale.create_engine')
    def test_close(self, mock_create_engine):
        mock_engine = Mock()
        mock_create_engine.return_value = mock_engine

        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        ts.connect()
        ts.close()

        mock_engine.dispose.assert_called_once()
        assert ts.engine is None

    def test_close_when_engine_none(self):
        ts = TimescaleToolKit(
            user='u', password='p', host='h', port=5432, database='d'
        )
        ts.close()
        assert ts.engine is None
