package com.tngtech.jgiven.testng;

import static java.util.Arrays.asList;

import java.lang.reflect.Method;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

import com.google.common.base.Throwables;
import com.tngtech.jgiven.exception.FailIfPassedException;
import org.testng.*;

import com.tngtech.jgiven.base.ScenarioTestBase;
import com.tngtech.jgiven.impl.ScenarioBase;
import com.tngtech.jgiven.impl.ScenarioHolder;
import com.tngtech.jgiven.impl.util.AssertionUtil;
import com.tngtech.jgiven.impl.util.ParameterNameUtil;
import com.tngtech.jgiven.report.impl.CommonReportHelper;
import com.tngtech.jgiven.report.model.NamedArgument;
import com.tngtech.jgiven.report.model.ReportModel;

/**
 * TestNG Test listener to enable JGiven for a test class
 */
public class ScenarioTestListener implements ITestListener {

    public static final String SCENARIO_ATTRIBUTE = "jgiven::scenario";
    public static final String REPORT_MODELS_ATTRIBUTE = "jgiven::reportModels";

    @Override
    public void onTestStart( ITestResult paramITestResult ) {
        Object instance = paramITestResult.getInstance();

        ScenarioBase scenario;

        if( instance instanceof ScenarioTestBase<?, ?, ?> ) {
            ScenarioTestBase<?, ?, ?> testInstance = (ScenarioTestBase<?, ?, ?>) instance;
            scenario = testInstance.createNewScenario();
        } else {
            scenario = new ScenarioBase();
        }

        ScenarioHolder.get().setScenarioOfCurrentThread( scenario );
        paramITestResult.setAttribute( SCENARIO_ATTRIBUTE, scenario );

        ReportModel reportModel = getReportModel( paramITestResult, instance.getClass() );
        scenario.setModel( reportModel );

        // TestNG does not work well when catching step exceptions, so we have to disable that feature
        // this mainly means that steps following a failing step are not reported in JGiven
        scenario.getExecutor().setSuppressStepExceptions(false);

        // avoid rethrowing exceptions as they are already thrown by the steps
        scenario.getExecutor().setSuppressExceptions(true);

        scenario.getExecutor().injectStages( instance );

        Method method = paramITestResult.getMethod().getConstructorOrMethod().getMethod();
        scenario.startScenario( instance.getClass(), method, getArgumentsFrom( method, paramITestResult ) );

        // inject state from the test itself
        scenario.getExecutor().readScenarioState( instance );
    }

    private ScenarioBase getScenario( ITestResult paramITestResult ) {
        return (ScenarioBase) paramITestResult.getAttribute( SCENARIO_ATTRIBUTE );
    }

    private ReportModel getReportModel(ITestResult testResult, Class<?> clazz) {
        ConcurrentHashMap<String, ReportModel> reportModels = getReportModels(testResult.getTestContext());

        ReportModel model = reportModels.get( clazz.getName() );
        if( model == null ) {
            model = new ReportModel();
            model.setTestClass( clazz );
            ReportModel previousModel = reportModels.putIfAbsent( clazz.getName(), model );
            if( previousModel != null ) {
                model = previousModel;
            }
        }
        AssertionUtil.assertNotNull( model, "Report model is null" );
        return model;
    }

    @Override
    public void onTestSuccess( ITestResult paramITestResult ) {
        testFinished( paramITestResult );
    }

    @Override
    public void onTestFailure( ITestResult paramITestResult ) {
        ScenarioBase scenario = getScenario( paramITestResult );
        if( scenario != null ) {
            scenario.getExecutor().failed( paramITestResult.getThrowable() );
            testFinished( paramITestResult );
        }
    }

    @Override
    public void onTestSkipped( ITestResult testResult ) {}

    private void testFinished( ITestResult testResult ) {
        try {
            ScenarioBase scenario = getScenario(testResult);
            scenario.finished();
        } catch( FailIfPassedException ex ) {
            testResult.setStatus(ITestResult.FAILURE);
            testResult.setThrowable(ex);
            testResult.getTestContext().getPassedTests().removeResult(testResult);
            testResult.getTestContext().getFailedTests().addResult(testResult, testResult.getMethod());
        } catch( Throwable throwable ) {
            Throwables.propagate(throwable);
        } finally {
            ScenarioHolder.get().removeScenarioOfCurrentThread();
        }
    }

    @Override
    public void onTestFailedButWithinSuccessPercentage( ITestResult paramITestResult ) {}

    @Override
    public void onStart( ITestContext paramITestContext ) {
        paramITestContext.setAttribute(REPORT_MODELS_ATTRIBUTE, new ConcurrentHashMap<String, ReportModel>());
    }

    @Override
    public void onFinish( ITestContext paramITestContext ) {
        ConcurrentHashMap<String, ReportModel> reportModels = getReportModels(paramITestContext);
        for( ReportModel reportModel : reportModels.values() ) {
            new CommonReportHelper().finishReport( reportModel );
        }
    }

    private ConcurrentHashMap<String, ReportModel> getReportModels(ITestContext paramITestContext) {
        return (ConcurrentHashMap<String, ReportModel>)
                paramITestContext.getAttribute(REPORT_MODELS_ATTRIBUTE);
    }

    private List<NamedArgument> getArgumentsFrom( Method method, ITestResult paramITestResult ) {
        return ParameterNameUtil.mapArgumentsWithParameterNames( method, asList( paramITestResult.getParameters() ) );
    }
}

@@ -80,275 +80,107 @@
 ort 
-java.lang.reflect.Method;%0Aimport java.util.List;%0Aimport java.util.concurrent.ConcurrentHashMap;%0A%0Aimport com.google.common.base.Throwables;%0Aimport com.tngtech.jgiven.exception.FailIfPassedException;%0Aimport org.testng.*;%0A%0Aimport com.tngtech.jgiven.base.ScenarioTestBase
+com.tngtech.jgiven.base.ScenarioTestBase;%0Aimport com.tngtech.jgiven.exception.FailIfPassedException
 ;%0Aim
@@ -538,16 +538,215 @@
 rtModel;
+%0Aimport java.lang.reflect.Method;%0Aimport java.util.List;%0Aimport java.util.concurrent.ConcurrentHashMap;%0Aimport org.testng.ITestContext;%0Aimport org.testng.ITestListener;%0Aimport org.testng.ITestResult;
 %0A%0A/**%0A *
@@ -4354,32 +4354,40 @@
 od());%0A        %7D
+%0A       
  catch( Throwabl
@@ -4418,28 +4418,34 @@
     
-T
+t
 hrow
-ables.propagate
+ new RuntimeException
 (thr

remove deprecated method