#include "builtins/builtins.h"
#include "precedence.h"
#include "builtins/helpers.h"

using namespace goose;
using namespace goose::eir;
using namespace goose::parse;
using namespace goose::builtins;

// Dollar is a low level operator: it doesn't expect a value rhs operand,
// but just an identifier eir instruction.
// So we can't use the convenient operator wrapper rule.
//
// Also, it works both as a prefix operator to construct a TVar,
// and as an infix operator (with a decl or TEXpr as the left value) to
// construct a Decl of the form "<type> $FOO" or a TDecl of the form
// "<TExpr> $FOO" (for instance $T $FOO)
namespace
{
    Value BuildOrRetrieveTVar( const Context& c, LocationId locationId, StringId name, bool forwarding )
    {
        if( name == "_"_sid )
            return forwarding ? ToValue( TTVar( "_"_sid ) ) : ToValue( TVar( "_"_sid ) );

        // Template name bindings are stored with a "$$" prefix so they don't
        // collide with regular names.
        auto captureIdentity = AppendToVectorTerm( c.identity(),
            TERM( forwarding ? "$$$"_sid : "$$"_sid ), TERM( name ) );

        // Look up the TVar name. If it exists, then we're inside of a template function
        // body, refering to a bound TVar, so we just return the stored value.
        // Otherwise, we're in a template declaration and have to construct a new TVar.
        Term result;

        switch( c.env()->retrieveValue( captureIdentity, c.identity(), result ) )
        {
            case sema::Env::Status::Success:
                return *EIRToValue( result );

            case sema::Env::Status::AmbiguousMatch:
                DiagnosticsManager::GetInstance().emitErrorMessage( locationId,
                    format( "unexpected ambiguous match when resolving '${}'.", name ) );
                return PoisonValue();

            default:
                break;
        }

        // No bound value was found: construct a TVar or a TTVar.
        return forwarding ? ToValue( TTVar( name ) ) : ToValue( TVar( name ) );
    }

    bool ParsePrefixDollarOperator( Parser& p, LocationId locationId, uint32_t prec )
    {
        auto nameTerm = p.resolver()->consumeUnresolved();
        if( !nameTerm )
        {
            DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( locationId,
                "expected an identifier after '$'.", 0 );
            return false;
        }

        const auto* name = get_if< StringId >( &nameTerm->first );
        if( !name )
        {
            DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( locationId,
                "expected an identifier after '$'.", 0 );
            return false;
        }

        const auto& c = p.context();
        auto val = BuildOrRetrieveTVar( c, nameTerm->second, *name, false );

        auto loc = Location::CreateSpanningLocation( locationId, nameTerm->second );
        val.setLocationId( loc );
        p.pushValue( move( val ) );
        return true;
    }

    bool ParsePrefixDollarDollarOperator( Parser& p, LocationId locationId, uint32_t prec )
    {
        auto nameTerm = p.resolver()->consumeUnresolved();
        if( !nameTerm )
        {
            DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( locationId,
                "expected an identifier after '$$'.", 0 );
            return false;
        }

        const auto* name = get_if< StringId >( &nameTerm->first );
        if( !name )
        {
            DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( locationId,
                "expected an identifier after '$$'.", 0 );
            return false;
        }

        const auto& c = p.context();
        auto val = BuildOrRetrieveTVar( c, nameTerm->second, *name, true );

        auto loc = Location::CreateSpanningLocation( locationId, nameTerm->second );
        val.setLocationId( loc );
        p.pushValue( move( val ) );
        return true;
    }

    optional< uint32_t > InfixDollarPrecedence( const Parser& p )
    {
        const auto& leftVal = p.peekLastValue();
        if( !leftVal )
            return nullopt;

        auto nameTerm = p.resolver()->lookAheadUnresolved();
        if( !nameTerm )
            return nullopt;

        if( !holds_alternative< StringId >( nameTerm->first ) )
            return nullopt;

        if( !builtins::IsType( p.context(), *leftVal ) && !builtins::IsTExpr( *leftVal ) )
            return nullopt;

        // The infix dollar operator is a syntactic shortcut meant to behave like we are applying
        // a TExpr to a type or another TExpr (just like a regular decl), so it have the application precedence.
        return precedence::Application;
    }

    bool ParseInfixDollarOperator( Parser& p, LocationId locationId, uint32_t prec )
    {
        const auto& leftVal = p.peekLastValue();
        if( !leftVal )
            return false;

        if( !builtins::IsType( p.context(), *leftVal ) && !IsTExpr( *leftVal )  )
            return false;

        auto nameTerm = p.resolver()->consumeUnresolved();
        const auto* name = get_if< StringId >( &nameTerm->first );
        const auto& c = p.context();

        auto val = BuildOrRetrieveTVar( c, nameTerm->second, *name, false );

        // If the TVar was already bound, we may end up with something in val
        // that isn't a TVar, which is an error.
        if( !IsTVar( val ) )
        {
            DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( nameTerm->second,
                "expected an unbound identifier after '$'." );
            return false;
        }

        auto loc = Location::CreateSpanningLocation( leftVal->locationId(), nameTerm->second );

        auto typeOrTExpr = ValueToEIR( *p.popValue() );
        auto tdecl = builtins::BuildTDecl( c, move( typeOrTExpr ), *name );
        assert( tdecl );

        p.pushValue( ToValue( move( *tdecl ) ).setLocationId( loc ) );
        return true;
    }
}

namespace goose::builtins
{
    void SetupDollarOp( Env& e )
    {
        Rule r( ParsePrefixDollarOperator, InfixDollarPrecedence, ParseInfixDollarOperator );
        RegisterRule( e, "$"_sid, move( r ) );

        Rule r2( ParsePrefixDollarDollarOperator );
        RegisterRule( e, "$$"_sid, move( r2 ) );
    }
}
