How to take screenshot on failure for both web and andorid native app

My test suite contains of both the driver , appium as well as web , reason , my scenario is like enabling button on web then verify that layout is visible of mobile app , I am using allure report for test logs , but problem I am facing is that , failure is on mobile test cases but screenshots is taken for web , what are best practice to construct such framework and how can I take screen shots based on specific platform failure. I had crated Listener which is been called for every class and in my class I have all my test method.

package com.qa.Listner;

import io.qameta.allure.Attachment;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.testng.ITestContext;
import org.testng.ITestListener;
import org.testng.ITestResult;

import com.aventstack.extentreports.MediaEntityBuilder;
import com.aventstack.extentreports.Status;
import com.qa.mysoultab.Base.BaseConfiguration;


public class ListenersT extends BaseConfiguration implements ITestListener {
    BaseConfiguration b = new BaseConfiguration();

    public void onFinish(ITestContext arg0) {
        // TODO Auto-generated method stub

    }

    public void onStart(ITestContext arg0) {
        // TODO Auto-generated method stub

    }

    public void onTestFailedButWithinSuccessPercentage(ITestResult arg0) {
        // TODO Auto-generated method stub

    }

    public void onTestFailure(ITestResult result) {
        WebDriver driver = BaseConfiguration.wd;
        if (driver instanceof WebDriver) {
            takeScreenshot(driver);
            childTest.log(Status.FAIL, MediaEntityBuilder
                    .createScreenCaptureFromBase64String(BaseConfiguration.getScreenshot(result.getName())).build());
        }
    }

    @Attachment(value = "Page screenshot", type = "image/png")
    public byte[] takeScreenshot(WebDriver driver) {
        byte[] screenshotAs = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);
        return screenshotAs;
    }

    public void onTestSkipped(ITestResult arg0) {
        // TODO Auto-generated method stub
        WebDriver driver = BaseConfiguration.wd;
        if (driver instanceof WebDriver) {
            takeScreenshot(driver);
            childTest.log(Status.SKIP, MediaEntityBuilder
                    .createScreenCaptureFromBase64String(BaseConfiguration.getScreenshot(arg0.getName())).build());
        }
    }

    public void onTestStart(ITestResult arg0) {
        // TODO Auto-generated method stub

    }

    public void onTestSuccess(ITestResult arg0) {
        // TODO Auto-generated method stub

    }

}

Sample Test Case

import com.Listner.ListenersT;
import com.Base.BaseConfiguration;
import io.qameta.allure.Description;
import io.qameta.allure.Severity;
import io.qameta.allure.SeverityLevel;
import io.qameta.allure.Story;
import org.apache.commons.lang3.StringUtils;
import org.testng.annotations.*;

import java.lang.reflect.Method;

import static org.assertj.core.api.Assertions.assertThat;

@Listeners({ ListenersT.class })
public class testcase extends BaseConfiguration {

    @Test(priority = 1)
    @Severity(SeverityLevel.BLOCKER)
    @Description("Verify changes is reflected.")

    public void tc001(Method method) {
        childTest = parentTest.createNode(StringUtils.capitalize(method.getName()));
        basePage.goToSections(env.webURL() + "/email");
        emailPage.enableButton();
        assertThat(email.isDisplay()).isTrue();
    }

    
    @Test(enabled = true,priority = 2,dependsOnMethods = "tc001")
    @Severity(SeverityLevel.BLOCKER)
    @Description("Verify section is visible on app.")

    public void tc002(Method method) {
        childTest = parentTest.createNode(StringUtils.capitalize(method.getName()));
        
        actTokenVerification.clickAndEnterEmail();
        actTokenVerification.clickOnProceedButton();
        assertThat(emailAct.isDiplayed()).isTrue();
    }



    @BeforeClass(alwaysRun = true)
    public void beforeClass() {
        parentTest = extentReport.createTest(getClass().getSimpleName());
        parentTest.assignCategory("Email");
        System.out.println("========================================================================================");
        System.out.println(getClass().getSimpleName());
        System.out.println("========================================================================================");
    }

    @BeforeMethod(alwaysRun = true)
    public void before(Method method) {
        System.out.println("Test Scenario: " + StringUtils.capitalize(method.getName()).replace("_", " "));
    }

    @AfterMethod(alwaysRun = true)
    public void consoleError(Method method) {
    }

    @AfterClass(alwaysRun = true)
    public void afterClass() {
        extentReport.flush();
    }

}