# SPDX-License-Identifier: MPL-2.0
# Copyright (c) 2025 Daniel Schmidt

"""Multi-archive report generator for PyATS test results.

This module handles generation of HTML reports from multiple PyATS archives,
supporting different test types (API, D2D) and creating combined summaries.
"""

import asyncio
import json
import logging
import os
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, cast

import aiofiles  # type: ignore[import-untyped]

from nac_test.core.constants import (
    DEBUG_MODE,
    HTML_REPORTS_DIRNAME,
    PYATS_RESULTS_DIRNAME,
    SUMMARY_REPORT_FILENAME,
)
from nac_test.pyats_core.reporting.generator import ReportGenerator
from nac_test.pyats_core.reporting.templates import TEMPLATES_DIR, get_jinja_environment
from nac_test.pyats_core.reporting.types import ResultStatus
from nac_test.pyats_core.reporting.utils.archive_extractor import ArchiveExtractor
from nac_test.pyats_core.reporting.utils.archive_inspector import ArchiveInspector
from nac_test.utils import get_or_create_event_loop

logger = logging.getLogger(__name__)


class MultiArchiveReportGenerator:
    """Handles report generation for multiple PyATS archives.

    This class extracts multiple archives, delegates to ReportGenerator for
    individual report generation, and creates a combined summary when needed.
    It maintains single responsibility by focusing only on multi-archive coordination.

    Attributes:
        output_dir: Base output directory for all reports
        pyats_results_dir: Directory where archives will be extracted
        env: Jinja2 environment for template rendering
    """

    def __init__(self, output_dir: Path, minimal_reports: bool = False) -> None:
        """Initialize the multi-archive report generator.

        Args:
            output_dir: Base output directory for all operations
            minimal_reports: Only include command outputs for failed/errored tests
        """
        self.output_dir = output_dir
        self.pyats_results_dir = output_dir / PYATS_RESULTS_DIRNAME
        self.minimal_reports = minimal_reports

        # Initialize Jinja2 environment for combined summary
        self.env = get_jinja_environment(TEMPLATES_DIR)

    async def generate_reports_from_archives(
        self, archive_paths: list[Path]
    ) -> dict[str, Any]:
        """Generate reports from multiple PyATS archives.

        This is the main entry point that coordinates the entire process:
        1. Extracts each archive to its appropriate subdirectory
        2. Runs ReportGenerator on each extracted archive
        3. Collects PyATS statistics (but does NOT generate combined summary)

        The combined summary is now generated by CombinedReportGenerator which
        includes both PyATS and Robot results.

        Args:
            archive_paths: List of paths to PyATS archive files

        Returns:
            Dictionary containing:
                - status: 'success', 'partial', or 'failed'
                - results: Dict mapping archive type to generation results
                - pyats_stats: Dict of PyATS statistics by framework (API/D2D)
                - duration: Total time taken
        """
        start_time = datetime.now()

        if not archive_paths:
            logger.warning("No archive paths provided")
            return {
                "status": "failed",
                "results": {},
                "pyats_stats": None,
                "duration": 0,
            }

        # Clean and prepare results directory
        if self.pyats_results_dir.exists():
            shutil.rmtree(self.pyats_results_dir)
        self.pyats_results_dir.mkdir(parents=True)

        # Note: JSONL files are no longer cleaned up by ReportGenerator.
        # Cleanup happens after stats collection (see below)
        logger.debug(
            f"Processing {len(archive_paths)} archive(s) for report generation"
        )

        # Process each archive
        results: dict[str, dict[str, Any]] = {}
        tasks = []

        for archive_path in archive_paths:
            if not archive_path.exists():
                logger.warning(f"Archive not found: {archive_path}")
                continue

            archive_type = ArchiveInspector.get_archive_type(archive_path)
            task = self._process_single_archive(archive_type, archive_path)
            tasks.append((archive_type, task))

        # Execute all archive processing in parallel
        if tasks:
            task_results = await asyncio.gather(
                *[task for _, task in tasks], return_exceptions=True
            )

            # Map results back to archive types
            for idx, (archive_type, _) in enumerate(tasks):
                if isinstance(task_results[idx], Exception):
                    logger.error(
                        f"Failed to process {archive_type} archive: {task_results[idx]}"
                    )
                    results[archive_type] = {
                        "status": "error",
                        "error": str(task_results[idx]),
                    }
                else:
                    results[archive_type] = cast(dict[str, Any], task_results[idx])

        # Collect PyATS stats for combined dashboard
        pyats_stats = None
        successful_archives = [
            k for k, v in results.items() if v.get("status") == "success"
        ]

        if len(successful_archives) > 0:
            pyats_stats = await self._collect_pyats_stats(results)
            logger.debug(
                f"Collected PyATS stats from {len(successful_archives)} archive type(s)"
            )
        else:
            logger.warning("No successful archives to collect stats from")

        # Cleanup JSONL files after stats collection
        # Respects NAC_TEST_DEBUG and NAC_TEST_PYATS_KEEP_REPORT_DATA environment variables
        if not (DEBUG_MODE or os.environ.get("NAC_TEST_PYATS_KEEP_REPORT_DATA")):
            logger.info("Cleaning up JSONL files after stats collection")
            await self._cleanup_all_jsonl_files()
        else:
            reason = (
                "NAC_TEST_DEBUG" if DEBUG_MODE else "NAC_TEST_PYATS_KEEP_REPORT_DATA"
            )
            logger.info(f"Keeping JSONL files ({reason} is set)")

        # Determine overall status
        if not results:
            overall_status = "failed"
        elif all(r.get("status") == "success" for r in results.values()):
            overall_status = "success"
        elif any(r.get("status") == "success" for r in results.values()):
            overall_status = "partial"
        else:
            overall_status = "failed"

        return {
            "status": overall_status,
            "duration": (datetime.now() - start_time).total_seconds(),
            "results": results,
            "pyats_stats": pyats_stats,  # Return stats for CombinedReportGenerator
        }

    async def _process_single_archive(
        self, archive_type: str, archive_path: Path
    ) -> dict[str, Any]:
        """Process a single archive by extracting and generating reports.

        Args:
            archive_type: Type of archive ('api' or 'd2d')
            archive_path: Path to the archive file

        Returns:
            Result dictionary from ReportGenerator
        """
        logger.info(f"Processing {archive_type} archive: {archive_path.name}")

        # Create type-specific extraction directory
        extract_dir = self.pyats_results_dir / archive_type
        extract_dir.mkdir(parents=True, exist_ok=True)

        try:
            # Extract archive
            await self._extract_archive(archive_path, extract_dir)

            # Calculate type-specific output directory where temp files are located
            # This ensures ReportGenerator can find the JSONL temp files in the correct subdirectory
            type_output_dir = self.output_dir / archive_type

            # Run ReportGenerator on extracted contents
            generator = ReportGenerator(
                type_output_dir,
                extract_dir,
                minimal_reports=self.minimal_reports,
                archive_type=archive_type,
            )
            result = await generator.generate_all_reports()

            # Add archive info to result
            result["archive_path"] = str(archive_path)
            result["archive_type"] = archive_type
            result["report_dir"] = str(extract_dir / HTML_REPORTS_DIRNAME)

            return result

        except Exception as e:
            logger.error(f"Failed to process {archive_type} archive: {e}")
            return {
                "status": "error",
                "error": str(e),
                "archive_path": str(archive_path),
                "archive_type": archive_type,
            }

    async def _extract_archive(self, archive_path: Path, target_dir: Path) -> None:
        """Extract a PyATS archive to the target directory using ArchiveExtractor.

        This method uses ArchiveExtractor to handle:
        - Clearing previous results
        - Preserving HTML reports in previous archives
        - Proper error handling

        Args:
            archive_path: Path to the archive file
            target_dir: Directory to extract contents to
        """
        loop = get_or_create_event_loop()

        def extract() -> None:
            # The target_dir already contains the full path like output_dir/pyats_results/api
            # So we just need to get the relative path from output_dir
            target_subdir = str(target_dir.relative_to(self.output_dir))

            # Use ArchiveExtractor for proper extraction with HTML preservation
            extraction_path = ArchiveExtractor.extract_archive_to_directory(
                archive_path, self.output_dir, target_subdir
            )

            if not extraction_path:
                raise RuntimeError(f"Failed to extract archive {archive_path}")

        await loop.run_in_executor(None, extract)
        logger.debug(
            f"Extracted {archive_path.name} to {target_dir} with HTML report preservation"
        )

    async def _read_jsonl_summary(self, jsonl_path: Path) -> dict[str, Any]:
        """Read only the summary record from a JSONL file.

        JSONL files contain multiple JSON objects, one per line. This method
        reads line-by-line to find and return the summary record which contains
        the overall_status and test metadata.

        Args:
            jsonl_path: Path to the JSONL result file.

        Returns:
            Dictionary containing summary information with overall_status.

        Raises:
            Exception: If file cannot be read or summary record not found.
        """
        summary = {}
        metadata = {}

        try:
            async with aiofiles.open(jsonl_path) as f:
                async for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    try:
                        record = json.loads(line)
                        record_type = record.get("type")

                        if record_type == "metadata":
                            metadata = record
                        elif record_type == "summary":
                            summary = record
                            # Found summary, can stop reading
                            break

                    except json.JSONDecodeError:
                        # Skip malformed lines
                        continue

        except Exception as e:
            logger.error(f"Failed to read JSONL file {jsonl_path}: {e}")
            raise

        # Return summary with overall_status
        return {
            "test_id": metadata.get("test_id") or summary.get("test_id"),
            "overall_status": summary.get("overall_status", ResultStatus.SKIPPED.value),
            "metadata": summary.get("metadata", {}),
        }

    async def _collect_pyats_stats(
        self, results: dict[str, dict[str, Any]]
    ) -> dict[str, dict[str, Any]]:
        """Collect PyATS statistics for combined dashboard.

        Reads JSONL files once to extract statistics for API and D2D test types.
        Does NOT generate combined_summary.html (that's done by CombinedReportGenerator).

        Args:
            results: Dictionary mapping archive types to their results

        Returns:
            Dictionary mapping archive type to stats dict.
            Format: {"API": {...stats...}, "D2D": {...stats...}}
            Each stats dict contains: title, total_tests, passed_tests, failed_tests,
            skipped_tests, success_rate, report_path
        """
        try:
            # Prepare test type specific statistics
            test_type_stats = {}

            for archive_type, result in results.items():
                if result.get("status") != "success":
                    continue

                # Read JSONL files from the archive's html_report_data directory
                archive_dir = self.pyats_results_dir / archive_type
                json_dir = archive_dir / HTML_REPORTS_DIRNAME / "html_report_data"

                stats = {
                    "title": "PyATS API"
                    if archive_type == "api"
                    else "PyATS Direct-to-Device (D2D)",
                    "total_tests": 0,
                    "passed_tests": 0,
                    "failed_tests": 0,
                    "skipped_tests": 0,
                    "success_rate": 0.0,
                    "report_path": f"{PYATS_RESULTS_DIRNAME}/{archive_type}/{HTML_REPORTS_DIRNAME}/{SUMMARY_REPORT_FILENAME}",
                }

                # Read all JSONL files to calculate statistics
                if json_dir.exists():
                    for jsonl_file in json_dir.glob("*.jsonl"):
                        try:
                            # Use helper method to read summary from JSONL
                            test_data = await self._read_jsonl_summary(jsonl_file)
                            status = test_data.get(
                                "overall_status", ResultStatus.SKIPPED.value
                            )

                            stats["total_tests"] = int(stats.get("total_tests", 0)) + 1

                            if status == ResultStatus.PASSED.value:
                                stats["passed_tests"] = (
                                    int(stats.get("passed_tests", 0)) + 1
                                )
                            elif status in [
                                ResultStatus.FAILED.value,
                                ResultStatus.ERRORED.value,
                            ]:
                                stats["failed_tests"] = (
                                    int(stats.get("failed_tests", 0)) + 1
                                )
                            elif status == ResultStatus.SKIPPED.value:
                                stats["skipped_tests"] = (
                                    int(stats.get("skipped_tests", 0)) + 1
                                )

                        except Exception as e:
                            logger.warning(
                                f"Failed to read test data from {jsonl_file}: {e}"
                            )

                # Calculate success rate for this test type
                total_tests = int(stats.get("total_tests", 0))
                skipped_tests = int(stats.get("skipped_tests", 0))
                passed_tests = int(stats.get("passed_tests", 0))

                tests_with_results = total_tests - skipped_tests
                if tests_with_results > 0:
                    stats["success_rate"] = (passed_tests / tests_with_results) * 100

                test_type_stats[archive_type.upper()] = stats

            logger.info(
                f"Collected PyATS stats for {len(test_type_stats)} archive type(s)"
            )
            return test_type_stats

        except Exception as e:
            logger.error(f"Failed to collect PyATS stats: {e}")
            return {}

    async def _cleanup_all_jsonl_files(self) -> None:
        """Clean up JSONL files from all archive directories after stats collection.

        This method is called after PyATS statistics have been collected to remove
        intermediate JSONL data files. Only HTML reports are kept for end users.

        The cleanup preserves disk space while maintaining all generated HTML reports
        which are the final user-facing output.
        """
        jsonl_count = 0
        try:
            for archive_dir in self.pyats_results_dir.iterdir():
                if not archive_dir.is_dir():
                    continue

                # Look for html_report_data directory
                json_dir = archive_dir / HTML_REPORTS_DIRNAME / "html_report_data"
                if not json_dir.exists():
                    continue

                # Remove all JSONL files
                jsonl_files = list(json_dir.glob("*.jsonl"))
                if jsonl_files:
                    logger.debug(
                        f"Cleaning up {len(jsonl_files)} JSONL file(s) from {archive_dir.name}"
                    )

                for jsonl_file in jsonl_files:
                    try:
                        jsonl_file.unlink()
                        jsonl_count += 1
                    except Exception as e:
                        logger.warning(f"Failed to delete {jsonl_file}: {e}")

            if jsonl_count > 0:
                logger.info(f"Cleaned up {jsonl_count} JSONL file(s) total")

        except Exception as e:
            logger.warning(f"Failed to cleanup JSONL files: {e}")
