"""
Integration tests for Task.run method using mocked S3 operations.
Tests the core workflow without requiring actual cloud credentials.
"""
import os
import sys
import tempfile
import shutil
import unittest
import sqlite3
from unittest.mock import patch, MagicMock, mock_open
from datetime import datetime, timezone
import pickle

# Add project root to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from o11ydiagnosis.subtasks.slowlogconverter.task import Task
from o11ydiagnosis.subtasks.slowlogconverter.slowlog import SlowLog, StreamParser
from o11ydiagnosis.subtasks.slowlogconverter.objstore import S3


class TestTaskIntegration(unittest.TestCase):
    """Integration tests for Task.run method"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.temp_dir = tempfile.mkdtemp()
        self.mock_s3 = MagicMock(spec=S3)
        
        # Create Task instance
        self.task = Task(
            s3=self.mock_s3,
            tenant="test-tenant",
            cluster="test-cluster",
            o11y_cluster="test-o11y-cluster",
            vendor="local",
            region="us-west-2",
            o11y_bucket="test-o11y-bucket",
            o11y_prefix="test/o11y/prefix",
            drhouse_bucket="test-drhouse-bucket",
            drhouse_prefix="test/drhouse/prefix",
            drhouse_local_path=self.temp_dir,
        )
        
        # Mock metrics
        self.metrics_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.metrics')
        self.mock_metrics = self.metrics_patcher.start()
        self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = None
        
        # Mock logger
        self.logger_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.log')
        self.mock_log = self.logger_patcher.start()
        
        # Mock write functions (task.py uses write_slowlogs_to_deltalake, not write_slowlogs_to_deltalake_multitables)
        self.write_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.write_slowlogs_to_deltalake')
        self.mock_write = self.write_patcher.start()
        
        # Mock optimize function
        self.optimize_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.try_to_optimize_table')
        self.mock_optimize = self.optimize_patcher.start()
    
    def tearDown(self):
        """Clean up test fixtures"""
        self.metrics_patcher.stop()
        self.logger_patcher.stop()
        self.write_patcher.stop()
        self.optimize_patcher.stop()
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
    
    def create_sample_slowlog_file(self, filename: str, num_logs: int = 1):
        """Create a sample slowlog file for testing"""
        file_path = os.path.join(self.temp_dir, filename)
        
        # Create sample slowlog content
        log_lines = []
        for i in range(num_logs):
            timestamp = datetime(2021, 7, 1, 10, 0, i, tzinfo=timezone.utc)
            log_lines.append(f"2021-07-01T10:00:{i:02d}.000000Z;test-instance;tidb;db-tidb-0;slowlog")
            log_lines.append("# Time: " + timestamp.strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
            log_lines.append("# Txn_start_ts: 0")
            log_lines.append("# User@Host: root[root] @ 127.0.0.1 [127.0.0.1]")
            log_lines.append("# Conn_ID: 12345")
            log_lines.append("# Query_time: 1.0")
            log_lines.append("# Digest: test_digest_123")
            log_lines.append("# Is_internal: false")
            log_lines.append(f"SELECT * FROM test_table_{i};")
            log_lines.append("")
        
        with open(file_path, 'w') as f:
            f.write('\n'.join(log_lines))
        
        return file_path
    
    def test_run_with_no_matched_keys(self):
        """Test run when no keys are matched"""
        hour_ts = 1625097600  # 2021-07-01 10:00:00 UTC
        
        # Mock find_matched_keys to return empty list
        with patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') as mock_find:
            mock_find.return_value = []
            
            # Create meta.db
            metadb_path = os.path.join(self.temp_dir, "meta.db")
            with sqlite3.connect(metadb_path) as metadb:
                metadb.isolation_level = None
                metadb.execute("""
                    CREATE TABLE IF NOT EXISTS "processed_objects" (
                        "timestamp" INTEGER NOT NULL,
                        "object_key" TEXT NOT NULL,
                        PRIMARY KEY ("timestamp", "object_key")
                    );
                """)
            
            # Run task
            self.task.run(hour_ts)
            
            # Verify metrics were updated
            self.mock_metrics.record_logs_written.assert_called_once()
            call_args = self.mock_metrics.record_logs_written.call_args[0]
            self.assertEqual(call_args[0], 0)  # 0 logs written
            self.assertEqual(call_args[1], "test-cluster")
    
    def test_run_with_new_keys(self):
        """Test run with new keys to process"""
        hour_ts = 1625097600  # 2021-07-01 10:00:00 UTC
        begin = hour_ts - hour_ts % 3600
        end = begin + 3600
        
        # Create sample log file
        log_filename = f"{int(begin)}_slowlog.log"
        self.create_sample_slowlog_file(log_filename, num_logs=2)
        
        # Mock S3 operations
        s3_key = f"test/o11y/prefix/logs/slowlog/{log_filename}"
        self.mock_s3.lists.return_value = [s3_key]
        self.mock_s3.existed.return_value = False  # No checkpoint exists
        self.mock_s3.download.return_value = True
        
        # Mock file operations to return our test file
        def mock_download(bucket, key, local_path):
            # Copy our test file to the download location
            src = os.path.join(self.temp_dir, log_filename)
            if os.path.exists(src):
                import shutil
                shutil.copy(src, local_path)
            return True
        
        self.mock_s3.download.side_effect = mock_download
        
        # Create meta.db
        metadb_path = os.path.join(self.temp_dir, "meta.db")
        with sqlite3.connect(metadb_path) as metadb:
            metadb.isolation_level = None
            metadb.execute("""
                CREATE TABLE IF NOT EXISTS "processed_objects" (
                    "timestamp" INTEGER NOT NULL,
                    "object_key" TEXT NOT NULL,
                    PRIMARY KEY ("timestamp", "object_key")
                );
            """)
        
        # Mock find_matched_keys to return our test key
        with patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') as mock_find:
            mock_find.return_value = [s3_key]
            
            # Mock deltalake_base_path to return local path
            with patch('o11ydiagnosis.subtasks.slowlogconverter.task.deltalake_base_path') as mock_base_path:
                mock_base_path.return_value = self.temp_dir
                
                # Run task
                self.task.run(hour_ts)
                
                # Verify S3 download was called
                self.mock_s3.download.assert_called()
                
                # Verify write function was called (if logs were parsed successfully)
                # Note: This depends on successful parsing, which may require more setup
    
    def test_run_with_existing_keys(self):
        """Test run when keys already exist in processed_objects"""
        hour_ts = 1625097600
        begin = hour_ts - hour_ts % 3600
        
        # Create meta.db with existing processed key
        metadb_path = os.path.join(self.temp_dir, "meta.db")
        with sqlite3.connect(metadb_path) as metadb:
            metadb.isolation_level = None
            metadb.execute("""
                CREATE TABLE IF NOT EXISTS "processed_objects" (
                    "timestamp" INTEGER NOT NULL,
                    "object_key" TEXT NOT NULL,
                    PRIMARY KEY ("timestamp", "object_key")
                );
            """)
            # Insert existing key
            metadb.execute(
                "INSERT INTO processed_objects (timestamp, object_key) VALUES (?, ?)",
                (begin, "existing_key.log")
            )
        
        # Mock find_matched_keys to return the existing key
        with patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') as mock_find:
            mock_find.return_value = ["existing_key.log"]
            
            # Mock deltalake_base_path
            with patch('o11ydiagnosis.subtasks.slowlogconverter.task.deltalake_base_path') as mock_base_path:
                mock_base_path.return_value = self.temp_dir
                
                # Run task
                self.task.run(hour_ts)
                
                # Verify no download occurred (key already processed)
                # Verify metrics were updated with 0 logs
                self.mock_metrics.record_logs_written.assert_called_once()
                call_args = self.mock_metrics.record_logs_written.call_args[0]
                self.assertEqual(call_args[0], 0)
    
    def test_run_temp_dir_creation(self):
        """Test that temp directory is created and cleaned"""
        hour_ts = 1625097600
        
        # Verify temp_dir exists after initialization
        self.assertTrue(os.path.exists(self.task.temp_dir))
        
        # Mock find_matched_keys to return empty list
        with patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') as mock_find:
            mock_find.return_value = []
            
            metadb_path = os.path.join(self.temp_dir, "meta.db")
            with sqlite3.connect(metadb_path) as metadb:
                metadb.isolation_level = None
                metadb.execute("""
                    CREATE TABLE IF NOT EXISTS "processed_objects" (
                        "timestamp" INTEGER NOT NULL,
                        "object_key" TEXT NOT NULL,
                        PRIMARY KEY ("timestamp", "object_key")
                    );
                """)
            
            self.task.run(hour_ts)
            
            # Temp dir should still exist
            self.assertTrue(os.path.exists(self.task.temp_dir))
    
    def test_run_optimize_check(self):
        """Test that optimize check is performed"""
        hour_ts = 1625097600
        
        # Mock checkpoint exists
        self.mock_s3.existed.return_value = True
        
        # Mock metrics to return old timestamp (should trigger optimize)
        self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = 0
        
        # Mock find_matched_keys to return empty list
        with patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') as mock_find:
            mock_find.return_value = []
            
            # Mock deltalake_base_path
            with patch('o11ydiagnosis.subtasks.slowlogconverter.task.deltalake_base_path') as mock_base_path:
                mock_base_path.return_value = self.temp_dir
                
                metadb_path = os.path.join(self.temp_dir, "meta.db")
                with sqlite3.connect(metadb_path) as metadb:
                    metadb.isolation_level = None
                    metadb.execute("""
                        CREATE TABLE IF NOT EXISTS "processed_objects" (
                            "timestamp" INTEGER NOT NULL,
                            "object_key" TEXT NOT NULL,
                            PRIMARY KEY ("timestamp", "object_key")
                        );
                    """)
                
                # Run task
                self.task.run(hour_ts)
                
                # Verify optimize was attempted (if checkpoint exists)
                # Note: This depends on the optimize logic in the actual code


if __name__ == '__main__':
    unittest.main()