"""
Test cases for o11ydiagnosis/subtasks/logreporter/task.py
"""
import os
import sys
import tempfile
import shutil
import unittest
import json
import sqlite3
from unittest.mock import patch, MagicMock, Mock, mock_open
from datetime import datetime

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

from o11ydiagnosis.subtasks.logreporter.task import Task, run_cluster
from o11ydiagnosis.subtasks.slowlogconverter.objstore import Local


class TestTask(unittest.TestCase):
    """Test Task class"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.temp_dir = tempfile.mkdtemp()
        self.local_path = os.path.join(self.temp_dir, "local")
        os.makedirs(self.local_path, exist_ok=True)
        
        self.mock_s3 = MagicMock()
        self.cluster = "test-cluster"
        self.bucket = "test-bucket"
        self.prefix = "test-prefix"
        
        self.task = Task(
            self.mock_s3,
            self.cluster,
            self.bucket,
            self.prefix,
            self.local_path
        )
    
    def tearDown(self):
        """Clean up test fixtures"""
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
    
    def test_task_initialization(self):
        """Test Task initialization"""
        self.assertEqual(self.task.s3, self.mock_s3)
        self.assertEqual(self.task.cluster, self.cluster)
        self.assertEqual(self.task.drhouse_bucket, self.bucket)
        self.assertEqual(self.task.drhouse_prefix, self.prefix)
        self.assertEqual(self.task.drhouse_local_path, self.local_path)
        self.assertEqual(self.task.existed_logdb_paths, {})
        self.assertEqual(self.task.daily_data_json_file_paths, {})
        self.assertEqual(self.task.benchmark_data_json_file_paths, {})
        self.assertEqual(self.task.daily_report_html_file_paths, {})
        self.assertEqual(self.task.all_statistic_file_paths, {})
    
    def test_task_initialization_strips_trailing_slashes(self):
        """Test that trailing slashes are stripped"""
        task = Task(
            self.mock_s3,
            self.cluster,
            self.bucket,
            self.prefix + "/",
            self.local_path + "/"
        )
        
        self.assertEqual(task.drhouse_prefix, self.prefix)
        self.assertEqual(task.drhouse_local_path, self.local_path)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_downloads_logdb_files(self, mock_log):
        """Test prepare downloads logdb files"""
        report_hour_ts = 1609459200  # 2021-01-01 00:00:00 UTC
        
        # Mock download to return True (file exists)
        self.mock_s3.download.return_value = True
        
        self.task.prepare(7, report_hour_ts)
        
        # Should attempt to download: 24 hour files + 7 benchmark files + 1 all_statistic file = 32
        self.assertEqual(self.mock_s3.download.call_count, 32)
        
        # Check that paths are tracked
        self.assertGreater(len(self.task.existed_logdb_paths), 0)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_uses_existing_logdb_files(self, mock_log):
        """Test prepare uses existing local logdb files"""
        report_hour_ts = 1609459200
        
        # Create a local file
        utc_hour = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d%H")
        logdb_local_path = os.path.join(
            self.local_path, "log-stats", "1h", f"{utc_hour}.db"
        )
        os.makedirs(os.path.dirname(logdb_local_path), exist_ok=True)
        with open(logdb_local_path, 'w') as f:
            f.write("test")
        
        self.task.prepare(7, report_hour_ts)
        
        # Should not download if file exists
        # But still should be tracked
        self.assertIn(logdb_local_path, self.task.existed_logdb_paths)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_downloads_benchmark_files(self, mock_log):
        """Test prepare downloads benchmark files"""
        report_hour_ts = 1609459200
        
        self.mock_s3.download.return_value = True
        
        self.task.prepare(7, report_hour_ts)
        
        # Should attempt to download benchmark files (7 days)
        # Plus 24 hour files, plus all_statistic file
        # Total calls should be at least 7 for benchmark
        download_calls = self.mock_s3.download.call_count
        self.assertGreaterEqual(download_calls, 7)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_downloads_all_statistic_file(self, mock_log):
        """Test prepare downloads all statistic file"""
        report_hour_ts = 1609459200
        
        self.mock_s3.download.return_value = True
        
        self.task.prepare(7, report_hour_ts)
        
        # Check that all_statistic file path is tracked
        all_statistic_path = os.path.join(
            self.local_path, "log-report", "all", f"{self.cluster}.csv"
        )
        # May or may not exist depending on download, but should be tracked if downloaded
        # The important thing is that the method was called
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.update_statistic_csv')
    @patch('o11ydiagnosis.subtasks.logreporter.task.send_feishu_message')
    @patch('o11ydiagnosis.subtasks.logreporter.task.generate_daily_statistic')
    @patch('o11ydiagnosis.subtasks.logreporter.task.generate_report')
    @patch('o11ydiagnosis.subtasks.logreporter.task.process_eigenvalue')
    @patch('o11ydiagnosis.subtasks.logreporter.task.process_daily_summary')
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_run_processes_data(self, mock_log, mock_process_summary, 
                                 mock_process_eigen, mock_gen_report, 
                                 mock_gen_stat, mock_send_feishu, mock_update_csv):
        """Test run processes data correctly"""
        report_hour_ts = 1609459200  # 2021-01-01 00:00:00 UTC
        
        # Setup mock return values
        mock_process_summary.return_value = [
            {"component": "comp1", "level": "ERROR", "total": 100}
        ]
        mock_process_eigen.return_value = [
            {"component": "comp1", "level": "ERROR", "total": 100, "tags": []}
        ]
        mock_gen_stat.return_value = {
            "all_logs_statistics": {"ERROR": {"comp1": 100}},
            "new_logs_statistics": {"new_error": 0, "new_warn": 0}
        }
        
        # Create some logdb files for process_daily_summary
        utc_hour = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d%H")
        logdb_path = os.path.join(
            self.local_path, "log-stats", "1h", f"{utc_hour}.db"
        )
        os.makedirs(os.path.dirname(logdb_path), exist_ok=True)
        conn = sqlite3.connect(logdb_path)
        cursor = conn.cursor()
        cursor.execute('''CREATE TABLE IF NOT EXISTS log_stats
                         (component TEXT, level TEXT, position TEXT, 
                          sample TEXT, count INTEGER)''')
        cursor.execute('''INSERT INTO log_stats VALUES (?, ?, ?, ?, ?)''',
                      ("comp1", "ERROR", "pos1", "Error message", 100))
        conn.commit()
        conn.close()
        self.task.existed_logdb_paths[logdb_path] = "s3-key"
        
        # Ensure directories exist
        utc_day = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d")
        json_dir = os.path.join(self.local_path, "log-report", "1d")
        html_dir = os.path.join(self.local_path, "log-report", "1d")
        csv_dir = os.path.join(self.local_path, "log-report", "all")
        os.makedirs(json_dir, exist_ok=True)
        os.makedirs(html_dir, exist_ok=True)
        os.makedirs(csv_dir, exist_ok=True)
        
        self.task.run(7, report_hour_ts)
        
        # Verify functions were called
        mock_process_summary.assert_called_once()
        mock_process_eigen.assert_called_once()
        mock_gen_report.assert_called_once()
        mock_gen_stat.assert_called_once()
        mock_send_feishu.assert_called_once()
        mock_update_csv.assert_called_once()
        
        # Check that files are tracked
        self.assertGreater(len(self.task.daily_data_json_file_paths), 0)
        self.assertGreater(len(self.task.daily_report_html_file_paths), 0)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_run_creates_daily_summary_json(self, mock_log):
        """Test run creates daily summary JSON file"""
        report_hour_ts = 1609459200
        
        # Create logdb file
        utc_hour = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d%H")
        logdb_path = os.path.join(
            self.local_path, "log-stats", "1h", f"{utc_hour}.db"
        )
        os.makedirs(os.path.dirname(logdb_path), exist_ok=True)
        conn = sqlite3.connect(logdb_path)
        cursor = conn.cursor()
        cursor.execute('''CREATE TABLE IF NOT EXISTS log_stats
                         (component TEXT, level TEXT, position TEXT, 
                          sample TEXT, count INTEGER)''')
        cursor.execute('''INSERT INTO log_stats VALUES (?, ?, ?, ?, ?)''',
                      ("comp1", "ERROR", "pos1", "Error", 100))
        conn.commit()
        conn.close()
        self.task.existed_logdb_paths[logdb_path] = "s3-key"
        
        # Ensure directory exists
        utc_day = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d")
        json_dir = os.path.join(self.local_path, "log-report", "1d")
        os.makedirs(json_dir, exist_ok=True)
        
        self.task.run(7, report_hour_ts)
        
        # Check JSON file was created
        json_path = os.path.join(json_dir, f"{utc_day}.json")
        self.assertTrue(os.path.exists(json_path))
        
        # Verify content
        with open(json_path, 'r') as f:
            data = json.load(f)
            self.assertIsInstance(data, list)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_run_creates_html_report(self, mock_log):
        """Test run creates HTML report"""
        report_hour_ts = 1609459200
        
        # Create logdb file
        utc_hour = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d%H")
        logdb_path = os.path.join(
            self.local_path, "log-stats", "1h", f"{utc_hour}.db"
        )
        os.makedirs(os.path.dirname(logdb_path), exist_ok=True)
        conn = sqlite3.connect(logdb_path)
        cursor = conn.cursor()
        cursor.execute('''CREATE TABLE IF NOT EXISTS log_stats
                         (component TEXT, level TEXT, position TEXT, 
                          sample TEXT, count INTEGER)''')
        cursor.execute('''INSERT INTO log_stats VALUES (?, ?, ?, ?, ?)''',
                      ("comp1", "ERROR", "pos1", "Error", 100))
        conn.commit()
        conn.close()
        self.task.existed_logdb_paths[logdb_path] = "s3-key"
        
        # Ensure directory exists
        utc_day = datetime.utcfromtimestamp(report_hour_ts).strftime("%Y%m%d")
        html_dir = os.path.join(self.local_path, "log-report", "1d")
        os.makedirs(html_dir, exist_ok=True)
        
        self.task.run(7, report_hour_ts)
        
        # Check HTML file was created
        html_path = os.path.join(html_dir, f"{utc_day}.html")
        self.assertTrue(os.path.exists(html_path))
        
        # Verify it's HTML
        with open(html_path, 'r') as f:
            content = f.read()
            self.assertIn("<!DOCTYPE html", content)
    
    def test_upload_modified_objects(self):
        """Test upload_modified_objects uploads all tracked files"""
        # Add some files to track
        test_file1 = os.path.join(self.temp_dir, "file1.json")
        test_file2 = os.path.join(self.temp_dir, "file2.html")
        test_file3 = os.path.join(self.temp_dir, "file3.csv")
        
        os.makedirs(os.path.dirname(test_file1), exist_ok=True)
        with open(test_file1, 'w') as f:
            f.write("test1")
        with open(test_file2, 'w') as f:
            f.write("test2")
        with open(test_file3, 'w') as f:
            f.write("test3")
        
        self.task.daily_data_json_file_paths[test_file1] = "s3-key1"
        self.task.daily_report_html_file_paths[test_file2] = "s3-key2"
        self.task.all_statistic_file_paths[test_file3] = "s3-key3"
        
        self.task.upload_modified_objects()
        
        # Verify upload was called for each file
        self.assertEqual(self.mock_s3.upload.call_count, 3)
        
        # Verify dictionaries are cleared
        self.assertEqual(len(self.task.daily_data_json_file_paths), 0)
        self.assertEqual(len(self.task.daily_report_html_file_paths), 0)
        self.assertEqual(len(self.task.all_statistic_file_paths), 0)
    
    def test_upload_modified_objects_empty(self):
        """Test upload_modified_objects with no files"""
        self.task.upload_modified_objects()
        
        # Should not call upload
        self.mock_s3.upload.assert_not_called()
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_asserts_timestamp(self, mock_log):
        """Test prepare asserts timestamp is at day boundary"""
        # Timestamp not at day boundary (has remainder)
        invalid_ts = 1609459200 + 3600  # Not divisible by 24*3600
        
        with self.assertRaises(AssertionError):
            self.task.prepare(7, invalid_ts)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_run_asserts_timestamp(self, mock_log):
        """Test run asserts timestamp is at day boundary"""
        # Timestamp not at day boundary
        invalid_ts = 1609459200 + 3600
        
        with self.assertRaises(AssertionError):
            self.task.run(7, invalid_ts)
    
    @patch('o11ydiagnosis.subtasks.logreporter.task.log')
    def test_prepare_handles_missing_s3_files(self, mock_log):
        """Test prepare handles missing S3 files gracefully"""
        report_hour_ts = 1609459200
        
        # Mock download to return False (file doesn't exist)
        self.mock_s3.download.return_value = False
        
        # Should not raise exception
        self.task.prepare(7, report_hour_ts)
        
        # Should still attempt downloads
        self.assertGreater(self.mock_s3.download.call_count, 0)


class TestRunCluster(unittest.TestCase):
    """Test run_cluster function"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.temp_dir = tempfile.mkdtemp()
        self.mock_s3 = MagicMock()
        self.task = Task(
            self.mock_s3,
            "test-cluster",
            "test-bucket",
            "test-prefix",
            self.temp_dir
        )
        # Import and reset n_thread
        from o11ydiagnosis.subtasks.logreporter import task as task_module
        task_module.n_thread = 0
    
    def tearDown(self):
        """Clean up test fixtures"""
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
        # Reset n_thread
        from o11ydiagnosis.subtasks.logreporter import task as task_module
        task_module.n_thread = 0
    
    def test_run_cluster_calls_task_methods(self):
        """Test run_cluster calls task methods"""
        ts = 1609459200  # Day boundary
        
        # Mock task methods
        self.task.prepare = MagicMock()
        self.task.run = MagicMock()
        self.task.upload_modified_objects = MagicMock()
        
        # Set n_thread before calling
        from o11ydiagnosis.subtasks.logreporter import task as task_module
        task_module.n_thread = 1
        
        run_cluster(self.task, ts)
        
        self.task.prepare.assert_called_once_with(7, ts)
        self.task.run.assert_called_once_with(7, ts)
        self.task.upload_modified_objects.assert_called_once()
        
        # n_thread should be decremented
        self.assertEqual(task_module.n_thread, 0)
    
    def test_run_cluster_handles_exceptions(self):
        """Test run_cluster handles exceptions gracefully"""
        ts = 1609459200
        
        # Set n_thread before calling
        from o11ydiagnosis.subtasks.logreporter import task as task_module
        task_module.n_thread = 1
        
        # Make prepare raise an exception
        original_prepare = self.task.prepare
        original_run = self.task.run
        original_upload = self.task.upload_modified_objects
        
        self.task.prepare = MagicMock(side_effect=Exception("Test error"))
        self.task.run = MagicMock()
        self.task.upload_modified_objects = MagicMock()
        
        # Should not raise exception (it's caught in the try-except-finally)
        # The exception is caught but not re-raised, so we can continue
        try:
            run_cluster(self.task, ts)
        except Exception as e:
            # If an exception is raised, it should only be from our test setup
            # The run_cluster function itself should handle exceptions
            if "Test error" not in str(e):
                raise
        
        # prepare should have been called
        self.task.prepare.assert_called_once()
        # n_thread should still be decremented in finally block
        self.assertEqual(task_module.n_thread, 0)
        # run and upload should not be called if prepare fails
        self.task.run.assert_not_called()
        self.task.upload_modified_objects.assert_not_called()
        
        # Restore original methods
        self.task.prepare = original_prepare
        self.task.run = original_run
        self.task.upload_modified_objects = original_upload


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