plugin.py

import json
import logging
import os
import shutil
import subprocess
import sys
import webbrowser
from pathlib import Path

import pytest

try:
    import tomllib
except ModuleNotFoundError:  # pragma: no cover - Python < 3.11
    import tomli as tomllib

from pytest_html_plus.compute_report_metadata import write_plus_metadata_if_main_worker
from pytest_html_plus.extract_link import extract_links_from_item
from pytest_html_plus.generate_html_report import JSONReporter
from pytest_html_plus.json_merge import merge_json_reports
from pytest_html_plus.json_to_xml_converter import convert_json_to_junit_xml
from pytest_html_plus.resolver_driver import resolve_driver, take_screenshot_generic
from pytest_html_plus.send_email_report import EmailSender
from pytest_html_plus.utils import (
    extract_error_block,
    extract_trace_block,
    load_email_env,
)

python_executable = shutil.which("python3") or shutil.which("python")
test_screenshot_paths = {}
PROFILE_OPTION = "--plus-profile"
PROFILE_SECTION = ("tool", "pytest-html-plus", "profiles")
PROFILE_OPTION_MAP = {
    "json-report": {"flag": "--json-report", "kind": "value"},
    "capture-screenshots": {"flag": "--capture-screenshots", "kind": "value"},
    "html-output": {"flag": "--html-output", "kind": "value"},
    "screenshots": {"flag": "--screenshots", "kind": "value"},
    "plus-email": {"flag": "--plus-email", "kind": "bool"},
    "should-open-report": {"flag": "--should-open-report", "kind": "value"},
    "generate-xml": {"flag": "--generate-xml", "kind": "bool"},
    "xml-report": {"flag": "--xml-report", "kind": "value"},
    "git-branch": {"flag": "--git-branch", "kind": "value"},
    "git-commit": {"flag": "--git-commit", "kind": "value"},
    "rp-env": {"flag": "--rp-env", "kind": "value"},
}


logger = logging.getLogger()
logger.setLevel(logging.INFO)

if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter("%(levelname)s - %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)


def _normalize_profile_key(key):
    return key.replace("_", "-")


def _find_pyproject_toml(start_path=None):
    current = Path(start_path or Path.cwd()).resolve()
    for candidate in (current, *current.parents):
        pyproject_file = candidate / "pyproject.toml"
        if pyproject_file.exists():
            return pyproject_file
    return None


def _read_profiles_from_pyproject(start_path=None):
    pyproject_file = _find_pyproject_toml(start_path=start_path)
    if pyproject_file is None:
        raise pytest.UsageError(
            f"{PROFILE_OPTION} requires a pyproject.toml file in the current "
            "directory or a parent directory"
        )

    try:
        with pyproject_file.open("rb") as fh:
            pyproject_data = tomllib.load(fh)
    except tomllib.TOMLDecodeError as exc:
        raise pytest.UsageError(f"Invalid TOML in {pyproject_file}: {exc}") from exc

    profiles = pyproject_data
    for section in PROFILE_SECTION:
        profiles = profiles.get(section, {})

    if not isinstance(profiles, dict):
        raise pytest.UsageError(
            f"Expected [{'.'.join(PROFILE_SECTION)}] to be a TOML table"
        )

    return profiles


def _build_profile_args(profile_name, start_path=None):
    profiles = _read_profiles_from_pyproject(start_path=start_path)
    profile = profiles.get(profile_name)

    if profile is None:
        available_profiles = ", ".join(sorted(profiles)) or "none"
        raise pytest.UsageError(
            f"Unknown profile '{profile_name}' for {PROFILE_OPTION}. "
            f"Available profiles: {available_profiles}"
        )

    if not isinstance(profile, dict):
        raise pytest.UsageError(
            f"Profile '{profile_name}' must be defined as a TOML table"
        )

    profile_args = []
    invalid_keys = []

    for raw_key, value in profile.items():
        key = _normalize_profile_key(raw_key)
        option = PROFILE_OPTION_MAP.get(key)

        if option is None:
            invalid_keys.append(raw_key)
            continue

        if option["kind"] == "bool":
            if not isinstance(value, bool):
                raise pytest.UsageError(
                    f"Profile '{profile_name}' option '{raw_key}' must be true or false"
                )
            if value:
                profile_args.append(option["flag"])
            continue

        if not isinstance(value, (str, int, float)):
            raise pytest.UsageError(
                f"Profile '{profile_name}' option '{raw_key}' "
                "must be a string-like value"
            )

        profile_args.append(f"{option['flag']}={value}")

    if invalid_keys:
        valid_keys = ", ".join(sorted(PROFILE_OPTION_MAP))
        invalid = ", ".join(sorted(invalid_keys))
        raise pytest.UsageError(
            f"Invalid key(s) in profile '{profile_name}': {invalid}. "
            f"Valid keys: {valid_keys}"
        )

    return profile_args


def apply_plus_profile_args(args, start_path=None):
    profile_name = None
    remaining_args = []
    skip_next = False

    for index, arg in enumerate(args):
        if skip_next:
            skip_next = False
            continue

        if arg == PROFILE_OPTION:
            if profile_name is not None:
                raise pytest.UsageError(f"{PROFILE_OPTION} can only be specified once")
            if index + 1 >= len(args):
                raise pytest.UsageError(f"{PROFILE_OPTION} requires a profile name")
            profile_name = args[index + 1]
            skip_next = True
            continue

        if arg.startswith(f"{PROFILE_OPTION}="):
            if profile_name is not None:
                raise pytest.UsageError(f"{PROFILE_OPTION} can only be specified once")
            profile_name = arg.split("=", 1)[1]
            if not profile_name:
                raise pytest.UsageError(f"{PROFILE_OPTION} requires a profile name")
            continue

        remaining_args.append(arg)

    if not profile_name:
        return list(args)

    profile_args = _build_profile_args(profile_name, start_path=start_path)
    return profile_args + remaining_args


@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item):
    if "caplog" not in item.fixturenames:
        item.fixturenames.append("caplog")


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(item, call):
    outcome = yield
    report = outcome.get_result()

    reporter = getattr(item.config, "_json_reporter", None)
    if not reporter:
        return
    error = None
    trace = None

    if (
        report.when == "call"
        or (report.when == "setup" and report.skipped)
        or (report.when in ("setup", "teardown") and report.failed)
    ):
        full_error = str(report.longrepr)
        error = extract_error_block(error=full_error)
        trace = extract_trace_block(full_error)

    if (
        report.when == "call"
        or (report.when == "setup" and report.skipped)
        or (report.when in ("setup", "teardown") and report.failed)
    ):
        config = item.config
        capture_option = config.getoption("--capture-screenshots")

        caplog_text = None
        if report.when in ("call", "setup"):
            if "caplog" in item.funcargs:
                caplog = item.funcargs["caplog"]
                try:
                    caplog_text = (
                        "\n".join(caplog.messages) if caplog.messages else None
                    )
                except KeyError:
                    caplog_text = None

        screenshot_path = config.getoption("--screenshots") or "screenshots"

        should_capture_screenshot = report.when in ("setup", "call") and (
            capture_option == "all"
            or (capture_option == "failed" and report.outcome == "failed")
        )

        if should_capture_screenshot:
            driver = resolve_driver(item)
            if driver:
                screenshot_path = take_screenshot_generic(screenshot_path, item, driver)

        worker_id = os.getenv("PYTEST_XDIST_WORKER") or "main"
        test_name = "".join(c if c.isalnum() else "_" for c in item.name)

        status = report.outcome
        if report.when in ("setup", "teardown") and report.failed:
            status = "error"

        reporter.log_result(
            test_name=test_name,
            nodeid=item.nodeid,
            status=status,
            duration=report.duration,
            attempt=None,
            error=error if report.failed else None,
            trace=trace if report.failed else None,
            markers=[m.name for m in item.iter_markers()],
            filepath=item.location[0],
            lineno=item.location[1],
            stdout=getattr(report, "capstdout", ""),
            stderr=getattr(report, "capstderr", ""),
            screenshot=screenshot_path,
            logs=caplog_text,
            worker=worker_id,
            links=extract_links_from_item(item),
        )


def pytest_sessionfinish(session, exitstatus):
    reporter = session.config._json_reporter

    raw_json_report = session.config.getoption("--json-report")
    html_output = session.config.getoption("--html-output") or "report_output"
    screenshots_path = session.config.getoption("--screenshots") or "screenshots"
    raw_xml_report = session.config.getoption("--xml-report")

    # ---- XML filename validation ----
    if raw_xml_report:
        if os.path.basename(raw_xml_report) != raw_xml_report:
            raise pytest.UsageError("--xml-report must be a filename, not a path")
        xml_filename = raw_xml_report
    else:
        xml_filename = "final_xml.xml"

    xml_path = os.path.join(html_output, xml_filename)

    os.makedirs(html_output, exist_ok=True)

    # ---- JSON filename validation ----
    if raw_json_report:
        if os.path.basename(raw_json_report) != raw_json_report:
            raise pytest.UsageError("--json-report must be a filename, not a path")
        json_filename = raw_json_report
    else:
        json_filename = "final_report.json"

    json_path = os.path.join(html_output, json_filename)
    reporter.report_path = json_path

    is_worker = os.getenv("PYTEST_XDIST_WORKER") is not None
    try:
        is_xdist = bool(session.config.getoption("-n"))
    except ValueError:
        is_xdist = False

    # ---- Worker behavior ----
    if is_worker:
        worker_id = os.getenv("PYTEST_XDIST_WORKER")
        worker_dir = ".pytest_worker_jsons"
        os.makedirs(worker_dir, exist_ok=True)

        reporter.report_path = os.path.join(worker_dir, f"{worker_id}.json")

        reporter.write_report()
        return

    # ---- Controller behavior ----
    # Always write raw results first
    reporter.write_report()

    # Always run merge (even for single worker)
    merge_json_reports(
        directory=".pytest_worker_jsons" if is_xdist else html_output,
        output_path=json_path,
    )

    script_path = os.path.join(os.path.dirname(__file__), "generate_html_report.py")
    if not os.path.exists(script_path):
        logger.warning(
            f"Report generation script not found at {script_path}. "
            f"Skipping HTML report generation."
        )
        return

    try:
        subprocess.run(
            [
                sys.executable,
                script_path,
                "--report",
                json_path,
                "--screenshots",
                screenshots_path,
                "--output",
                html_output,
            ],
            check=True,
        )
    except Exception as e:
        raise RuntimeError(f"Exception during HTML report generation: {e}") from e

    # ---- Generate XML ----
    if session.config.getoption("--generate-xml"):
        try:
            convert_json_to_junit_xml(json_path, xml_path)
            print(f"XML report generated: {xml_path}")
        except Exception as e:
            raise RuntimeError(f"Failed to generate XML report: {e}") from e

    if not os.getenv("PYTEST_XDIST_WORKER"):
        if os.path.exists(screenshots_path):
            try:
                shutil.rmtree(screenshots_path)
            except Exception:
                logger.warning("Could not clean up screenshots directory")

    if session.config.getoption("--plus-email"):
        try:
            config = load_email_env()
            config["report_path"] = html_output
            sender = EmailSender(config, report_path=html_output)
            sender.send()
        except Exception as e:
            raise RuntimeError(f"Failed to send email: {e}") from e

    # ---- Open report (controller only) ----
    open_html_report(
        report_path=os.path.join(html_output, "report.html"),
        json_path=json_path,
        config=session.config,
    )


def pytest_sessionstart(session):
    html_output = session.config.getoption("--html-output") or "report_output"
    git_branch = (
        session.config.getoption("--git-branch")
        or "Pass --git-branch to populate git metadata"
    )
    git_commit = (
        session.config.getoption("--git-commit")
        or "Pass --git-commit to populate git metadata"
    )
    rp_env = (
        session.config.getoption("--rp-env")
        or "Pass --rp-env <name> to populate environment"
    )
    configure_logging()
    session.config.addinivalue_line(
        "markers", "link(url): Add a link to external test case or documentation."
    )
    write_plus_metadata_if_main_worker(
        session.config,
        report_path=html_output,
        git_branch=git_branch,
        git_commit=git_commit,
        rp_env=rp_env,
    )


def pytest_load_initial_conftests(args):
    args[:] = apply_plus_profile_args(args, start_path=Path.cwd())
    if not any(arg.startswith("--capture") for arg in args):
        args.append("--capture=tee-sys")


def pytest_addoption(parser):
    group = parser.getgroup("pytest-html-plus", "pytest-html-plus reporting options")

    group.addoption(
        PROFILE_OPTION,
        action="store",
        default=None,
        help="Load pytest-html-plus options from a named profile in pyproject.toml",
    )
    group.addoption(
        "--json-report",
        action="store",
        default="final_report.json",
        help="Name of the JSON report file generated alongside the HTML report",
    )
    group.addoption(
        "--capture-screenshots",
        action="store",
        default="failed",
        choices=["failed", "all", "none"],
        help="Capture screenshots: failed (default), all, or none",
    )
    group.addoption("--html-output", default="report_output")
    group.addoption("--screenshots", default="screenshots")
    group.addoption(
        "--plus-email",
        action="store_true",
        default=False,
        help="Send HTML test report via email after test run",
    )
    group.addoption(
        "--should-open-report",
        action="store",
        default="failed",
        choices=["always", "failed", "never"],
        help="When to open the HTML report: always, failed, or never (default: failed)",
    )
    group.addoption(
        "--generate-xml",
        action="store_true",
        default=False,
        help="Generate JUnit-style XML from the final JSON report",
    )
    group.addoption(
        "--xml-report",
        action="store",
        default=None,
        help="Name of the XML report file generated alongside the HTML report (used with --generate-xml)",  # noqa
    )
    group.addoption(
        "--git-branch",
        action="store",
        default="Pass --git-branch to populate git metadata",
        help="Helps show branch information on the report",
    )
    group.addoption(
        "--git-commit",
        action="store",
        default="Pass --git-commit to populate git metadata",
        help="Helps show commitId information on the report",
    )
    group.addoption(
        "--rp-env",
        action="store",
        default="Pass --rp-env to populate environment",
        help="Helps show env information on the report",
    )


def configure_logging():
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if not any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
        handler = logging.StreamHandler(sys.stdout)
        handler.setLevel(logging.INFO)
        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
        handler.setFormatter(formatter)
        logger.addHandler(handler)


def pytest_configure(config):
    global _saved_config
    _saved_config = config

    INTERNAL_JSON_DIR = Path(".pytest_worker_jsons")
    report_path = config.getoption("--json-report") or "final_report.json"
    worker_id = os.getenv("PYTEST_XDIST_WORKER")

    if worker_id:
        INTERNAL_JSON_DIR.mkdir(parents=True, exist_ok=True)
        name, ext = os.path.splitext(report_path)
        report_path = INTERNAL_JSON_DIR / f"{name}_{worker_id}{ext}"

    reporter = JSONReporter(report_path=report_path)
    config._json_reporter = reporter
    config.attempt_counters = {}


def pytest_collectreport(report):
    if report.failed:
        global _saved_config
        reporter = getattr(_saved_config, "_json_reporter", None)
        if reporter:
            reporter.log_result(
                test_name="COLLECTION ERROR",
                nodeid=str(report.nodeid),
                status="error",
                duration=getattr(report, "duration", 0.0),
                error=str(report.longrepr),
                markers=[],
                filepath=str(report.fspath),
                lineno=0,
                stdout="",
                stderr="",
                screenshot=None,
                logs=None,
                worker=os.getenv("PYTEST_XDIST_WORKER") or "main",
            )


def mark_flaky_tests(results):
    # Group test attempts by nodeid
    tests_by_nodeid = {}
    for test in results:
        tests_by_nodeid.setdefault(test["nodeid"], []).append(test)

    # Only return the final test attempt with flaky info
    final_results = []
    for nodeid, attempts in tests_by_nodeid.items():
        final_test = attempts[-1].copy()

        previous_statuses = [t["status"] for t in attempts[:-1]]
        final_status = final_test["status"]

        # A test is flaky if it passed at the end but had at least one failure before
        if final_status == "passed" and "failed" in previous_statuses:
            final_test["flaky"] = True
            final_test["flaky_attempts"] = [t["status"] for t in attempts]
        else:
            final_test["flaky"] = False

        final_results.append(final_test)

    return final_results


def open_html_report(report_path: str, json_path: str, config) -> None:
    if os.environ.get("CI") == "true":
        return

    should_open = config.getoption("--should-open-report", default="failed").lower()

    if not report_path or not os.path.exists(report_path):
        return

    try:
        with open(json_path, encoding="utf-8") as f:
            report_data = json.load(f)

        results = report_data.get("results", [])

        has_failures = any(
            t.get("status") == "failed" or t.get("error") for t in results
        )

        if should_open == "always" or (should_open == "failed" and has_failures):
            webbrowser.open(f"file://{os.path.abspath(report_path)}")

    except Exception as e:
        try:
            logger.warning(f"Could not open report in browser: {e}")
        except Exception:
            print(f"Could not open report in browser: {e}")