/*
 * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlinx.dataframe.plugin

import org.jetbrains.kotlin.fir.expressions.FirExpression
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter

internal inline fun <reified T> KotlinTypeFacade.analyzeRefinedCallShape(
    call: FirFunctionCall,
    expectedReturnType: ClassId,
    reporter: InterpretationErrorReporter,
): CallResult<T>? {
    val callReturnType = call.resolvedType
    if (callReturnType.classId != expectedReturnType) return null
    // rootMarker is expected to be a token generated by the plugin.
    // it's implied by "refined call"
    // thus ConeClassLikeType
    val rootMarkers = callReturnType.typeArguments.filterIsInstance<ConeClassLikeType>()
    if (rootMarkers.size != callReturnType.typeArguments.size) return null

    val newSchema: T? = call.interpreterName(session)?.let { name ->
        when (name) {
            else -> name.load<Interpreter<*>>(isTest)?.let { processor ->
                val dataFrameSchema = interpret(call, processor, reporter = reporter)
                    .let {
                        val value = it?.value
                        if (value !is T) {
                            if (!reporter.errorReported) {
                                reporter.reportInterpretationError(call, "${processor::class} must return ${T::class}, but was $value")
                            }
                            null
                        } else {
                            value
                        }
                    }
                dataFrameSchema
            }
        }
    }

    return CallResult(rootMarkers, newSchema)
}

data class CallResult<T>(val markers: List<ConeClassLikeType>, val result: T?)

class RefinedArguments(val refinedArguments: List<RefinedArgument>) : List<RefinedArgument> by refinedArguments

data class RefinedArgument(val name: Name, val expression: FirExpression) {

    override fun toString(): String {
        return "RefinedArgument(name=$name, expression=${expression})"
    }
}
